def test_change_state_running(self, mock_get_kube_client,
mock_kubernetes_job_watcher,
mock_kube_config):
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,
mock_kube_config):
executor = KubernetesExecutor()
executor.start()
key = ('dag_id', 'task_id', 'ex_time', 'try_number3')
executor._change_state(key, State.FAILED, 'pod_id', 'default')
self.assertTrue(executor.event_buffer[key] == State.FAILED)
mock_delete_pod.assert_called_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()
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_with('pod_id', 'default')
def test_change_state_skip_pod_deletion(self, mock_delete_pod, mock_get_kube_client,
mock_kubernetes_job_watcher, mock_kube_config):
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()
