def test_port(self): port = Port('http', 80) k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, ports=[port], ) context = self.create_context(k) k.execute(context=context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0]['ports'] = [{ 'name': 'http', 'containerPort': 80 }] self.assertEqual(self.expected_pod, actual_pod)
def test_image_pull_policy_not_set(self, mock_client, monitor_mock, start_mock): from airflow.utils.state import State k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, cluster_context='default', ) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context=context) self.assertEqual( start_mock.call_args[0][0].spec.containers[0].image_pull_policy, 'IfNotPresent', )
def test_pod_dnspolicy(self): dns_policy = "ClusterFirstWithHostNet" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test-" + str(random.randint(0, 1000000)), task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, hostnetwork=True, dnspolicy=dns_policy, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['hostNetwork'] = True self.expected_pod['spec']['dnsPolicy'] = dns_policy assert self.expected_pod['spec'] == actual_pod['spec'] assert self.expected_pod['metadata']['labels'] == actual_pod[ 'metadata']['labels']
def test_no_need_to_describe_pod_on_success(self, mock_client, monitor_mock, start_mock): from airflow.utils.state import State name_base = 'test' k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name=name_base, task_id="task", in_cluster=False, do_xcom_push=False, cluster_context='default', ) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context=context) assert not mock_client.return_value.read_namespaced_pod.called
def test_port(self): port = k8s.V1ContainerPort( name='http', container_port=80, ) k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test-" + str(random.randint(0, 1000000)), task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, ports=[port], ) context = create_context(k) k.execute(context=context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0]['ports'] = [{'name': 'http', 'containerPort': 80}] assert self.expected_pod == actual_pod
def test_randomize_pod_name(self, mock_client, monitor_mock, start_mock): from airflow.utils.state import State name_base = 'test' k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name=name_base, task_id="task", in_cluster=False, do_xcom_push=False, cluster_context='default', ) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context=context) assert start_mock.call_args[0][0].metadata.name.startswith(name_base) assert start_mock.call_args[0][0].metadata.name != name_base
def test_image_pull_secrets_correctly_set(self, mock_client, monitor_mock, start_mock): fake_pull_secrets = "fakeSecret" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, image_pull_secrets=fake_pull_secrets, cluster_context='default', ) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context=context) self.assertEqual( start_mock.call_args[0][0].spec.image_pull_secrets, [k8s.V1LocalObjectReference(name=fake_pull_secrets)], )
def test_envs_from_configmaps(self, mock_client, mock_monitor, mock_start): # GIVEN configmap = 'test-configmap' # WHEN k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, configmaps=[configmap], ) # THEN mock_monitor.return_value = (State.SUCCESS, None) context = create_context(k) k.execute(context) assert mock_start.call_args[0][0].spec.containers[0].env_from == [ k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap)) ]
def test_pod_priority_class_name(self, mock_client, monitor_mock, start_mock): # pylint: disable=unused-argument """Test ability to assign priorityClassName to pod""" priority_class_name = "medium-test" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, priority_class_name=priority_class_name, ) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['priorityClassName'] = priority_class_name self.assertEqual(self.expected_pod, actual_pod)
def test_pod_failure(self): """ Tests that the task fails when a pod reports a failure """ bad_internal_command = ["foobar 10 "] k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=bad_internal_command, labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, ) with self.assertRaises(AirflowException): context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0][ 'args'] = bad_internal_command self.assertEqual(self.expected_pod, actual_pod)
def test_run_as_user_root(self): security_context = { 'securityContext': { 'runAsUser': 0, } } k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, security_context=security_context, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['securityContext'] = security_context self.assertEqual(self.expected_pod, actual_pod)
def test_pod_resources(self): resources = { 'limit_cpu': 0.25, 'limit_memory': '64Mi', 'limit_ephemeral_storage': '2Gi', 'request_cpu': '250m', 'request_memory': '64Mi', 'request_ephemeral_storage': '1Gi', } k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, resources=resources, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0]['resources'] = { 'requests': { 'memory': '64Mi', 'cpu': '250m', 'ephemeral-storage': '1Gi' }, 'limits': { 'memory': '64Mi', 'cpu': 0.25, 'ephemeral-storage': '2Gi' }, } assert self.expected_pod == actual_pod
def test_full_pod_spec(self): pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, namespace="default", name="test-pod"), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="perl", command=["/bin/bash"], args=[ "-c", 'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json' ], env=[k8s.V1EnvVar(name="env_name", value="value")], ) ], restart_policy="Never", ), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, full_pod_spec=pod_spec, do_xcom_push=True, is_delete_operator_pod=False, ) context = create_context(k) result = k.execute(context) assert result is not None assert k.pod.metadata.labels == { 'fizz': 'buzz', 'foo': 'bar', 'airflow_version': mock.ANY, 'dag_id': 'dag', 'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b', 'kubernetes_pod_operator': 'True', 'task_id': mock.ANY, 'try_number': '1', } assert k.pod.spec.containers[0].env == [ k8s.V1EnvVar(name="env_name", value="value") ] assert result == {"hello": "world"}
def test_pod_delete_even_on_launcher_error( self, mock_client, delete_pod_mock, monitor_pod_mock, start_pod_mock): k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, cluster_context='default', is_delete_operator_pod=True, ) monitor_pod_mock.side_effect = AirflowException('fake failure') with self.assertRaises(AirflowException): context = self.create_context(k) k.execute(context=context) assert delete_pod_mock.called
def test_envs_from_secrets(self, mock_client, monitor_mock, start_mock): # GIVEN secret_ref = 'secret_name' secrets = [Secret('env', None, secret_ref)] # WHEN k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], secrets=secrets, labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, ) # THEN monitor_mock.return_value = (State.SUCCESS, None) context = create_context(k) k.execute(context) assert start_mock.call_args[0][0].spec.containers[0].env_from == [ k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name=secret_ref)) ]
def test_fs_group(self): security_context = { 'securityContext': { 'fsGroup': 1000, } } k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test-fs-group", task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, security_context=security_context, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['securityContext'] = security_context self.assertEqual(self.expected_pod, actual_pod)
def test_config_path_move(self): new_config_path = '/tmp/kube_config' old_config_path = get_kubeconfig_path() shutil.copy(old_config_path, new_config_path) k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test1", task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, is_delete_operator_pod=False, config_file=new_config_path, ) context = create_context(k) k.execute(context) expected_pod = copy(self.expected_pod) expected_pod['metadata']['labels']['already_checked'] = 'True' actual_pod = self.api_client.sanitize_for_serialization(k.pod) assert expected_pod == actual_pod
def test_envs_from_configmaps(self, mock_client, mock_monitor, mock_start): # GIVEN from airflow.utils.state import State configmap_name = "test-config-map" env_from = [k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap_name))] # WHEN k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test-" + str(random.randint(0, 1000000)), task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, env_from=env_from, ) # THEN mock_monitor.return_value = (State.SUCCESS, None) context = create_context(k) k.execute(context) self.assertEqual(mock_start.call_args[0][0].spec.containers[0].env_from, env_from)
def test_config_path(self, client_mock, launcher_mock): from airflow.utils.state import State file_path = "/tmp/fake_file" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, config_file=file_path, cluster_context='default', ) launcher_mock.return_value = (State.SUCCESS, None) k.execute(None) client_mock.assert_called_once_with( in_cluster=False, cluster_context='default', config_file=file_path, )
def test_config_path(self): file_path = "/tmp/fake_file" k = KubernetesPodOperator( namespace="default", image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=False, config_file=file_path, cluster_context="default", ) self.monitor_mock.return_value = (State.SUCCESS, None) self.client_mock.list_namespaced_pod.return_value = [] context = self.create_context(k) k.execute(context=context) self.client_mock.assert_called_once_with( in_cluster=False, cluster_context="default", config_file=file_path, )
def test_reattach_failing_pod_once(self): client = kube_client.get_kube_client(in_cluster=False) name = "test" namespace = "default" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["exit 1"], labels={"foo": "bar"}, name="test", task_id=name, in_cluster=False, do_xcom_push=False, is_delete_operator_pod=False, termination_grace_period=0, ) context = create_context(k) # launch pod with mock.patch( "airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion" ) as await_pod_completion_mock: pod_mock = MagicMock() # we don't want failure because we don't want the pod to be patched as "already_checked" pod_mock.status.phase = 'Succeeded' await_pod_completion_mock.return_value = pod_mock k.execute(context) name = k.pod.metadata.name pod = client.read_namespaced_pod(name=name, namespace=namespace) while pod.status.phase != "Failed": pod = client.read_namespaced_pod(name=name, namespace=namespace) assert 'already_checked' not in pod.metadata.labels # should not call `create_pod`, because there's a pod there it should find # should use the found pod and patch as "already_checked" (in failure block) with mock.patch( "airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod" ) as create_mock: with pytest.raises(AirflowException): k.execute(context) pod = client.read_namespaced_pod(name=name, namespace=namespace) assert pod.metadata.labels["already_checked"] == "True" create_mock.assert_not_called() # `create_pod` should be called because though there's still a pod to be found, # it will be `already_checked` with mock.patch( "airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod" ) as create_mock: with pytest.raises(AirflowException): k.execute(context) create_mock.assert_called_once()
def test_pod_template_file_with_overrides_system(self): fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml' k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), labels={"foo": "bar", "fizz": "buzz"}, env_vars=[k8s.V1EnvVar(name="env_name", value="value")], in_cluster=False, pod_template_file=fixture, do_xcom_push=True, ) context = create_context(k) result = k.execute(context) self.assertIsNotNone(result) self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'}) self.assertEqual(k.pod.spec.containers[0].env, [k8s.V1EnvVar(name="env_name", value="value")]) self.assertDictEqual(result, {"hello": "world"})
def test_pod_template_file_with_full_pod_spec(self): fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml' pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, ), spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", env=[k8s.V1EnvVar(name="env_name", value="value")], ) ]), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, pod_template_file=fixture, full_pod_spec=pod_spec, do_xcom_push=True, ) context = create_context(k) result = k.execute(context) assert result is not None assert k.pod.metadata.labels == { 'fizz': 'buzz', 'foo': 'bar', 'airflow_version': mock.ANY, 'dag_id': 'dag', 'execution_date': mock.ANY, 'kubernetes_pod_operator': 'True', 'task_id': mock.ANY, 'try_number': '1', } assert k.pod.spec.containers[0].env == [ k8s.V1EnvVar(name="env_name", value="value") ] assert result == {"hello": "world"}
def test_full_pod_spec(self): pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, namespace="default", name="test-pod"), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="perl", command=["/bin/bash"], args=[ "-c", 'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json' ], env=[k8s.V1EnvVar(name="env_name", value="value")], ) ], restart_policy="Never", ), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, full_pod_spec=pod_spec, do_xcom_push=True, ) context = create_context(k) result = k.execute(context) self.assertIsNotNone(result) self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'}) self.assertEqual(k.pod.spec.containers[0].env, [k8s.V1EnvVar(name="env_name", value="value")]) self.assertDictEqual(result, {"hello": "world"})
def test_xcom_push(self): return_value = '{"foo": "bar"\n, "buzz": 2}' args = ['echo \'{}\' > /airflow/xcom/return.json'.format(return_value)] k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=args, labels={"foo": "bar"}, name="test", task_id="task", in_cluster=False, do_xcom_push=True, ) self.assertEqual(k.execute(None), json.loads(return_value)) actual_pod = self.api_client.sanitize_for_serialization(k.pod) volume = self.api_client.sanitize_for_serialization(PodDefaults.VOLUME) volume_mount = self.api_client.sanitize_for_serialization(PodDefaults.VOLUME_MOUNT) container = self.api_client.sanitize_for_serialization(PodDefaults.SIDECAR_CONTAINER) self.expected_pod['spec']['containers'][0]['args'] = args self.expected_pod['spec']['containers'][0]['volumeMounts'].insert(0, volume_mount) self.expected_pod['spec']['volumes'].insert(0, volume) self.expected_pod['spec']['containers'].append(container) self.assertEqual(self.expected_pod, actual_pod)
def test_reattach_failing_pod_once(self): from airflow.utils.state import State client = kube_client.get_kube_client(in_cluster=False) name = "test" namespace = "default" k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["exit 1"], labels={"foo": "bar"}, name="test", task_id=name, in_cluster=False, do_xcom_push=False, is_delete_operator_pod=False, termination_grace_period=0, ) context = create_context(k) with mock.patch( "airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.monitor_pod" ) as monitor_mock: monitor_mock.return_value = (State.SUCCESS, None, None) k.execute(context) name = k.pod.metadata.name pod = client.read_namespaced_pod(name=name, namespace=namespace) while pod.status.phase != "Failed": pod = client.read_namespaced_pod(name=name, namespace=namespace) with pytest.raises(AirflowException): k.execute(context) pod = client.read_namespaced_pod(name=name, namespace=namespace) assert pod.metadata.labels["already_checked"] == "True" with mock.patch( "airflow.providers.cncf.kubernetes" ".operators.kubernetes_pod.KubernetesPodOperator" ".create_new_pod_for_operator" ) as create_mock: create_mock.return_value = ("success", {}, {}) k.execute(context) create_mock.assert_called_once()
def test_init_container(self): # GIVEN volume_mounts = [ k8s.V1VolumeMount(mount_path='/etc/foo', name='test-volume', sub_path=None, read_only=True) ] init_environments = [ k8s.V1EnvVar(name='key1', value='value1'), k8s.V1EnvVar(name='key2', value='value2'), ] init_container = k8s.V1Container( name="init-container", image="ubuntu:16.04", env=init_environments, volume_mounts=volume_mounts, command=["bash", "-cx"], args=["echo 10"], ) volume_config = {'persistentVolumeClaim': {'claimName': 'test-volume'}} volume = Volume(name='test-volume', configs=volume_config) expected_init_container = { 'name': 'init-container', 'image': 'ubuntu:16.04', 'command': ['bash', '-cx'], 'args': ['echo 10'], 'env': [{ 'name': 'key1', 'value': 'value1' }, { 'name': 'key2', 'value': 'value2' }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'test-volume', 'readOnly': True }], } k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test", task_id="task", volumes=[volume], init_containers=[init_container], in_cluster=False, do_xcom_push=False, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['initContainers'] = [expected_init_container] self.expected_pod['spec']['volumes'] = [{ 'name': 'test-volume', 'persistentVolumeClaim': { 'claimName': 'test-volume' } }] assert self.expected_pod == actual_pod
def test_pod_template_file( self, mock_client, monitor_mock, start_mock # pylint: disable=unused-argument ): from airflow.utils.state import State path = sys.path[0] + '/tests/kubernetes/pod.yaml' k = KubernetesPodOperator(task_id="task" + self.get_current_task_name(), pod_template_file=path, do_xcom_push=True) monitor_mock.return_value = (State.SUCCESS, None) context = create_context(k) with self.assertLogs(k.log, level=logging.DEBUG) as cm: k.execute(context) expected_line = textwrap.dedent("""\ DEBUG:airflow.task.operators:Starting pod: api_version: v1 kind: Pod metadata: annotations: {} cluster_name: null creation_timestamp: null deletion_grace_period_seconds: null\ """).strip() self.assertTrue( any(line.startswith(expected_line) for line in cm.output)) actual_pod = self.api_client.sanitize_for_serialization(k.pod) expected_dict = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': {}, 'labels': {}, 'name': 'memory-demo', 'namespace': 'mem-example' }, 'spec': { 'affinity': {}, 'containers': [ { 'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'], 'command': ['stress'], 'env': [], 'envFrom': [], 'image': 'apache/airflow:stress-2020.07.10-1.0.4', 'name': 'base', 'ports': [], 'resources': { 'limits': { 'memory': '200Mi' }, 'requests': { 'memory': '100Mi' } }, 'volumeMounts': [{ 'mountPath': '/airflow/xcom', 'name': 'xcom' }], }, { 'command': [ 'sh', '-c', 'trap "exit 0" INT; while true; do sleep 30; done;' ], 'image': 'alpine', 'name': 'airflow-xcom-sidecar', 'resources': { 'requests': { 'cpu': '1m' } }, 'volumeMounts': [{ 'mountPath': '/airflow/xcom', 'name': 'xcom' }], }, ], 'hostNetwork': False, 'imagePullSecrets': [], 'initContainers': [], 'restartPolicy': 'Never', 'securityContext': {}, 'serviceAccountName': 'default', 'tolerations': [], 'volumes': [{ 'emptyDir': {}, 'name': 'xcom' }], }, } self.assertEqual(expected_dict, actual_pod)
def test_pod_template_file(self, mock_client, monitor_mock, start_mock): # pylint: disable=unused-argument from airflow.utils.state import State k = KubernetesPodOperator( task_id='task', pod_template_file='tests/kubernetes/pod.yaml', do_xcom_push=True) monitor_mock.return_value = (State.SUCCESS, None) context = self.create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': ANY, 'namespace': 'mem-example' }, 'spec': { 'volumes': [{ 'name': 'xcom', 'emptyDir': {} }], 'containers': [{ 'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'], 'command': ['stress'], 'image': 'polinux/stress', 'name': 'memory-demo-ctr', 'resources': { 'limits': { 'memory': '200Mi' }, 'requests': { 'memory': '100Mi' } }, 'volumeMounts': [{ 'name': 'xcom', 'mountPath': '/airflow/xcom' }] }, { 'name': 'airflow-xcom-sidecar', 'image': "alpine", 'command': ['sh', '-c', PodDefaults.XCOM_CMD], 'volumeMounts': [{ 'name': 'xcom', 'mountPath': '/airflow/xcom' }], 'resources': { 'requests': { 'cpu': '1m' } }, }], } }, actual_pod)
def test_pod_template_file(self, mock_client, await_pod_completion_mock, create_mock, extract_xcom_mock): extract_xcom_mock.return_value = '{}' path = sys.path[0] + '/tests/kubernetes/pod.yaml' k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), random_name_suffix=False, pod_template_file=path, do_xcom_push=True, ) pod_mock = MagicMock() pod_mock.status.phase = 'Succeeded' await_pod_completion_mock.return_value = pod_mock context = create_context(k) with self.assertLogs(k.log, level=logging.DEBUG) as cm: k.execute(context) expected_line = textwrap.dedent("""\ DEBUG:airflow.task.operators:Starting pod: api_version: v1 kind: Pod metadata: annotations: {} cluster_name: null creation_timestamp: null deletion_grace_period_seconds: null\ """).strip() assert any(line.startswith(expected_line) for line in cm.output) actual_pod = self.api_client.sanitize_for_serialization(k.pod) expected_dict = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': {}, 'labels': { 'dag_id': 'dag', 'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b', 'kubernetes_pod_operator': 'True', 'task_id': mock.ANY, 'try_number': '1', }, 'name': 'memory-demo', 'namespace': 'mem-example', }, 'spec': { 'affinity': {}, 'containers': [ { 'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'], 'command': ['stress'], 'env': [], 'envFrom': [], 'image': 'ghcr.io/apache/airflow-stress:1.0.4-2021.07.04', 'name': 'base', 'ports': [], 'resources': { 'limits': { 'memory': '200Mi' }, 'requests': { 'memory': '100Mi' } }, 'volumeMounts': [{ 'mountPath': '/airflow/xcom', 'name': 'xcom' }], }, { 'command': [ 'sh', '-c', 'trap "exit 0" INT; while true; do sleep 1; done;' ], 'image': 'alpine', 'name': 'airflow-xcom-sidecar', 'resources': { 'requests': { 'cpu': '1m' } }, 'volumeMounts': [{ 'mountPath': '/airflow/xcom', 'name': 'xcom' }], }, ], 'hostNetwork': False, 'imagePullSecrets': [], 'initContainers': [], 'nodeSelector': {}, 'restartPolicy': 'Never', 'securityContext': {}, 'tolerations': [], 'volumes': [{ 'emptyDir': {}, 'name': 'xcom' }], }, } version = actual_pod['metadata']['labels']['airflow_version'] assert version.startswith(airflow_version) del actual_pod['metadata']['labels']['airflow_version'] assert expected_dict == actual_pod