def test_run_next_exception(self, mock_get_kube_client,
mock_kubernetes_job_watcher):
# When a quota is exceeded this is the ApiException we get
r = HTTPResponse()
r.body = {
"kind":
"Status",
"apiVersion":
"v1",
"metadata": {},
"status":
"Failure",
"message":
"pods \"podname\" is forbidden: " +
"exceeded quota: compute-resources, " +
"requested: limits.memory=4Gi, " + "used: limits.memory=6508Mi, " +
"limited: limits.memory=10Gi",
"reason":
"Forbidden",
"details": {
"name": "podname",
"kind": "pods"
},
"code":
403
},
r.status = 403
r.reason = "Forbidden"
# A mock kube_client that throws errors when making a pod
mock_kube_client = mock.patch('kubernetes.client.CoreV1Api',
autospec=True)
mock_kube_client.create_namespaced_pod = mock.MagicMock(
side_effect=ApiException(http_resp=r))
mock_get_kube_client.return_value = mock_kube_client
kubernetesExecutor = KubernetesExecutor()
kubernetesExecutor.start()
# Execute a task while the Api Throws errors
try_number = 1
kubernetesExecutor.execute_async(key=('dag', 'task', datetime.utcnow(),
try_number),
command='command',
executor_config={})
kubernetesExecutor.sync()
kubernetesExecutor.sync()
assert mock_kube_client.create_namespaced_pod.called
self.assertFalse(kubernetesExecutor.task_queue.empty())
# Disable the ApiException
mock_kube_client.create_namespaced_pod.side_effect = None
# Execute the task without errors should empty the queue
kubernetesExecutor.sync()
assert mock_kube_client.create_namespaced_pod.called
self.assertTrue(kubernetesExecutor.task_queue.empty())
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()
def test_run_next_exception(self, mock_get_kube_client, mock_kubernetes_job_watcher):
# When a quota is exceeded this is the ApiException we get
r = HTTPResponse()
r.body = {
"kind": "Status",
"apiVersion": "v1",
"metadata": {},
"status": "Failure",
"message": "pods \"podname\" is forbidden: " +
"exceeded quota: compute-resources, " +
"requested: limits.memory=4Gi, " +
"used: limits.memory=6508Mi, " +
"limited: limits.memory=10Gi",
"reason": "Forbidden",
"details": {"name": "podname", "kind": "pods"},
"code": 403},
r.status = 403
r.reason = "Forbidden"
# A mock kube_client that throws errors when making a pod
mock_kube_client = mock.patch('kubernetes.client.CoreV1Api', autospec=True)
mock_kube_client.create_namespaced_pod = mock.MagicMock(
side_effect=ApiException(http_resp=r))
mock_get_kube_client.return_value = mock_kube_client
kubernetesExecutor = KubernetesExecutor()
kubernetesExecutor.start()
# Execute a task while the Api Throws errors
try_number = 1
kubernetesExecutor.execute_async(key=('dag', 'task', datetime.utcnow(), try_number),
command='command', executor_config={})
kubernetesExecutor.sync()
kubernetesExecutor.sync()
mock_kube_client.create_namespaced_pod.assert_called()
self.assertFalse(kubernetesExecutor.task_queue.empty())
# Disable the ApiException
mock_kube_client.create_namespaced_pod.side_effect = None
# Execute the task without errors should empty the queue
kubernetesExecutor.sync()
mock_kube_client.create_namespaced_pod.assert_called()
self.assertTrue(kubernetesExecutor.task_queue.empty())
