def setUp(self):
self.watcher = KubernetesJobWatcher(
namespace="airflow",
multi_namespace_mode=False,
watcher_queue=mock.MagicMock(),
resource_version="0",
scheduler_job_id="123",
kube_config=mock.MagicMock(),
)
self.kube_client = mock.MagicMock()
self.core_annotations = {
"dag_id": "dag",
"task_id": "task",
"execution_date": "dt",
"try_number": "1",
}
self.pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="foo",
annotations={
"airflow-worker": "bar",
**self.core_annotations
},
namespace="airflow",
resource_version="456",
),
status=k8s.V1PodStatus(phase="Pending"),
)
self.events = []
def setUp(self) -> None:
self.watcher_queue = multiprocessing.Manager().Queue()
self.watcher = KubernetesJobWatcher(
namespace="namespace",
multi_namespace_mode=False,
watcher_queue=self.watcher_queue,
resource_version="0",
worker_uuid="0",
kube_config=None,
)
class TestKubernetesJobWatcher(unittest.TestCase):
def setUp(self) -> None:
self.watcher_queue = multiprocessing.Manager().Queue()
self.watcher = KubernetesJobWatcher(
namespace="namespace",
multi_namespace_mode=False,
watcher_queue=self.watcher_queue,
resource_version="0",
worker_uuid="0",
kube_config=None,
)
def test_running_task(self):
self.watcher.process_status(pod_id="pod_id",
namespace="namespace",
status="Running",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"})
self.assertTrue(self.watcher_queue.empty())
def test_succeeded_task(self):
self.watcher.process_status(pod_id="pod_id",
namespace="namespace",
status="Succeeded",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"})
result = self.watcher_queue.get_nowait()
self.assertEqual(('pod_id', 'namespace', None, {
'foo': 'bar'
}, '5'), result)
self.assertTrue(self.watcher_queue.empty())
def test_failed_task(self):
self.watcher.process_status(pod_id="pod_id",
namespace="namespace",
status="Failed",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"})
result = self.watcher_queue.get_nowait()
self.assertEqual(('pod_id', 'namespace', "failed", {
'foo': 'bar'
}, '5'), result)
self.assertTrue(self.watcher_queue.empty())
class TestKubernetesJobWatcher(unittest.TestCase):
def setUp(self):
self.watcher = KubernetesJobWatcher(
namespace="airflow",
multi_namespace_mode=False,
watcher_queue=mock.MagicMock(),
resource_version="0",
scheduler_job_id="123",
kube_config=mock.MagicMock(),
)
self.kube_client = mock.MagicMock()
self.core_annotations = {
"dag_id": "dag",
"task_id": "task",
"execution_date": "dt",
"try_number": "1",
}
self.pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="foo",
annotations={
"airflow-worker": "bar",
**self.core_annotations
},
namespace="airflow",
resource_version="456",
),
status=k8s.V1PodStatus(phase="Pending"),
)
self.events = []
def _run(self):
with mock.patch(
'airflow.executors.kubernetes_executor.watch') as mock_watch:
mock_watch.Watch.return_value.stream.return_value = self.events
latest_resource_version = self.watcher._run(
self.kube_client,
self.watcher.resource_version,
self.watcher.scheduler_job_id,
self.watcher.kube_config,
)
assert self.pod.metadata.resource_version == latest_resource_version
def assert_watcher_queue_called_once_with_state(self, state):
self.watcher.watcher_queue.put.assert_called_once_with((
self.pod.metadata.name,
self.watcher.namespace,
state,
self.core_annotations,
self.pod.metadata.resource_version,
))
def test_process_status_pending(self):
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()
def test_process_status_pending_deleted(self):
self.events.append({"type": 'DELETED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(State.FAILED)
def test_process_status_failed(self):
self.pod.status.phase = "Failed"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(State.FAILED)
def test_process_status_succeeded(self):
self.pod.status.phase = "Succeeded"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(None)
def test_process_status_running(self):
self.pod.status.phase = "Running"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()
def test_process_status_catchall(self):
self.pod.status.phase = "Unknown"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()
@mock.patch.object(KubernetesJobWatcher, 'process_error')
def test_process_error_event_for_410(self, mock_process_error):
message = "too old resource version: 27272 (43334)"
self.pod.status.phase = 'Pending'
self.pod.metadata.resource_version = '0'
mock_process_error.return_value = '0'
raw_object = {"code": 410, "message": message}
self.events.append({
"type": "ERROR",
"object": self.pod,
"raw_object": raw_object
})
self._run()
mock_process_error.assert_called_once_with(self.events[0])
def test_process_error_event_for_raise_if_not_410(self):
message = "Failure message"
self.pod.status.phase = 'Pending'
raw_object = {"code": 422, "message": message, "reason": "Test"}
self.events.append({
"type": "ERROR",
"object": self.pod,
"raw_object": raw_object
})
with self.assertRaises(AirflowException) as e:
self._run()
assert str(
e.exception
) == 'Kubernetes failure for {} with code {} and message: {}'.format(
raw_object['reason'],
raw_object['code'],
raw_object['message'],
)
class TestKubernetesJobWatcher(unittest.TestCase):
def setUp(self):
self.watcher = KubernetesJobWatcher(
namespace="airflow",
multi_namespace_mode=False,
watcher_queue=mock.MagicMock(),
resource_version="0",
scheduler_job_id="123",
kube_config=mock.MagicMock(),
)
self.kube_client = mock.MagicMock()
self.core_annotations = {
"dag_id": "dag",
"task_id": "task",
"execution_date": "dt",
"try_number": "1",
}
self.pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="foo",
annotations={
"airflow-worker": "bar",
**self.core_annotations
},
namespace="airflow",
resource_version="456",
),
status=k8s.V1PodStatus(phase="Pending"),
)
self.events = []
def _run(self):
with mock.patch(
'airflow.executors.kubernetes_executor.watch') as mock_watch:
mock_watch.Watch.return_value.stream.return_value = self.events
latest_resource_version = self.watcher._run(
self.kube_client,
self.watcher.resource_version,
self.watcher.scheduler_job_id,
self.watcher.kube_config,
)
assert self.pod.metadata.resource_version == latest_resource_version
def assert_watcher_queue_called_once_with_state(self, state):
self.watcher.watcher_queue.put.assert_called_once_with((
self.pod.metadata.name,
self.watcher.namespace,
state,
self.core_annotations,
self.pod.metadata.resource_version,
))
def test_process_status_pending(self):
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()
def test_process_status_pending_deleted(self):
self.events.append({"type": 'DELETED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(State.FAILED)
def test_process_status_failed(self):
self.pod.status.phase = "Failed"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(State.FAILED)
def test_process_status_succeeded(self):
self.pod.status.phase = "Succeeded"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.assert_watcher_queue_called_once_with_state(None)
def test_process_status_running(self):
self.pod.status.phase = "Running"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()
def test_process_status_catchall(self):
self.pod.status.phase = "Unknown"
self.events.append({"type": 'MODIFIED', "object": self.pod})
self._run()
self.watcher.watcher_queue.put.assert_not_called()