def test_change_state_running(self, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number1') executor._change_state(key, State.RUNNING, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.RUNNING)
def test_change_state_failed(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number3') executor._change_state(key, State.FAILED, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_called_once_with('pod_id')
def test_change_state_success(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() test_time = timezone.utcnow() key = ('dag_id', 'task_id', test_time, 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id', 'default') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_called_once_with('pod_id', 'default')
def test_change_state_success(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher, mock_kube_config): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_called_with('pod_id')
def test_change_state_skip_pod_deletion(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods = False executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_not_called()
def test_change_state_failed_pod_deletion(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods_on_failure = True executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.FAILED, 'pod_id', 'test-namespace') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_called_once_with('pod_id', 'test-namespace')
def test_change_state_failed(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods = False executor.kube_config.delete_worker_pods_on_failure = False executor.start() test_time = timezone.utcnow() key = ('dag_id', 'task_id', test_time, 'try_number3') executor._change_state(key, State.FAILED, 'pod_id', 'default') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_not_called()