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())