class TestPodGenerator(unittest.TestCase): def setUp(self): self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'} self.secrets = [ # This should be a secretRef Secret('env', None, 'secret_a'), # This should be a single secret mounted in volumeMounts Secret('volume', '/etc/foo', 'secret_b'), # This should produce a single secret mounted in env Secret('env', 'TARGET', 'secret_b', 'source_b'), ] self.resources = Resources('1Gi', 1, '2Gi', 2, 1) self.k8s_client = ApiClient() self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod-0', 'labels': { 'app': 'myapp' }, 'namespace': 'default' }, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'args': [], 'command': ['sh', '-c', 'echo Hello Kubernetes!'], 'imagePullPolicy': 'IfNotPresent', 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }, { 'secretRef': { 'name': 'secret_a' } }], 'resources': { 'requests': { 'memory': '1Gi', 'cpu': 1 }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 1 }, }, 'ports': [{ 'name': 'foo', 'containerPort': 1234 }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol0', 'readOnly': True }] }], 'restartPolicy': 'Never', 'volumes': [{ 'name': 'secretvol0', 'secret': { 'secretName': 'secret_b' } }], 'hostNetwork': False, 'imagePullSecrets': [{ 'name': 'pull_secret_a' }, { 'name': 'pull_secret_b' }], 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } } @mock.patch('uuid.uuid4') def test_gen_pod(self, mock_uuid): mock_uuid.return_value = '0' pod_generator = PodGenerator( labels={'app': 'myapp'}, name='myapp-pod', image_pull_secrets='pull_secret_a,pull_secret_b', image='busybox', envs=self.envs, cmds=['sh', '-c', 'echo Hello Kubernetes!'], security_context=k8s.V1PodSecurityContext( run_as_user=1000, fs_group=2000, ), namespace='default', ports=[k8s.V1ContainerPort(name='foo', container_port=1234)], configmaps=['configmap_a', 'configmap_b']) result = pod_generator.gen_pod() result = append_to_pod(result, self.secrets) result = self.resources.attach_to_pod(result) result_dict = self.k8s_client.sanitize_for_serialization(result) # sort result_dict['spec']['containers'][0]['env'].sort( key=lambda x: x['name']) result_dict['spec']['containers'][0]['envFrom'].sort( key=lambda x: list(x.values())[0]['name']) self.assertDictEqual(result_dict, self.expected) @mock.patch('uuid.uuid4') def test_gen_pod_extract_xcom(self, mock_uuid): mock_uuid.return_value = '0' pod_generator = PodGenerator( labels={'app': 'myapp'}, name='myapp-pod', image_pull_secrets='pull_secret_a,pull_secret_b', image='busybox', envs=self.envs, cmds=['sh', '-c', 'echo Hello Kubernetes!'], namespace='default', security_context=k8s.V1PodSecurityContext( run_as_user=1000, fs_group=2000, ), ports=[k8s.V1ContainerPort(name='foo', container_port=1234)], configmaps=['configmap_a', 'configmap_b']) pod_generator.extract_xcom = True result = pod_generator.gen_pod() result = append_to_pod(result, self.secrets) result = self.resources.attach_to_pod(result) result_dict = self.k8s_client.sanitize_for_serialization(result) container_two = { 'name': 'airflow-xcom-sidecar', 'image': "alpine", 'command': ['sh', '-c', PodDefaults.XCOM_CMD], 'volumeMounts': [{ 'name': 'xcom', 'mountPath': '/airflow/xcom' }] } self.expected['spec']['containers'].append(container_two) self.expected['spec']['containers'][0]['volumeMounts'].insert( 0, { 'name': 'xcom', 'mountPath': '/airflow/xcom' }) self.expected['spec']['volumes'].insert(0, { 'name': 'xcom', 'emptyDir': {} }) result_dict['spec']['containers'][0]['env'].sort( key=lambda x: x['name']) self.assertEqual(result_dict, self.expected) @mock.patch('uuid.uuid4') def test_from_obj(self, mock_uuid): mock_uuid.return_value = '0' result = PodGenerator.from_obj({ "KubernetesExecutor": { "annotations": { "test": "annotation" }, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": { "path": "/tmp/" }, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], "securityContext": { "runAsUser": 1000 } } }) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': { 'test': 'annotation' }, }, 'spec': { 'containers': [{ 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume' }], }], 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result) def test_reconcile_pods(self): with mock.patch('uuid.uuid4') as mock_uuid: mock_uuid.return_value = '0' base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=k8s.V1ContainerPort(name='port', container_port=2118), volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = PodGenerator(envs={ 'key2': 'val2' }, image='', name='name2', cmds=['/bin/command2.sh', 'arg2'], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' }]).gen_pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( result, { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'name2-0' }, 'spec': { 'containers': [{ 'args': [], 'command': ['/bin/command1.sh', 'arg1'], 'env': [{ 'name': 'key1', 'value': 'val1' }, { 'name': 'key2', 'value': 'val2' }], 'envFrom': [], 'image': 'image1', 'imagePullPolicy': 'IfNotPresent', 'name': 'base', 'ports': { 'containerPort': 2118, 'name': 'port', }, 'volumeMounts': [ { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }, { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' } ] }], 'hostNetwork': False, 'imagePullSecrets': [], 'restartPolicy': 'Never', 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }, { 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }] } })
class TestPodGenerator(unittest.TestCase): def setUp(self): self.static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48') self.deserialize_result = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'memory-demo', 'namespace': 'mem-example' }, 'spec': { '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' } } }] } } self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'} self.secrets = [ # This should be a secretRef Secret('env', None, 'secret_a'), # This should be a single secret mounted in volumeMounts Secret('volume', '/etc/foo', 'secret_b'), # This should produce a single secret mounted in env Secret('env', 'TARGET', 'secret_b', 'source_b'), ] self.labels = { 'airflow-worker': 'uuid', 'dag_id': 'dag_id', 'execution_date': 'date', 'task_id': 'task_id', 'try_number': '3', 'airflow_version': mock.ANY, 'kubernetes_executor': 'True' } self.metadata = { 'labels': self.labels, 'name': 'pod_id-' + self.static_uuid.hex, 'namespace': 'namespace' } self.resources = Resources('1Gi', 1, '2Gi', '2Gi', 2, 1, '4Gi') self.k8s_client = ApiClient() self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod-' + self.static_uuid.hex, 'labels': { 'app': 'myapp' }, 'namespace': 'default' }, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'args': [], 'command': ['sh', '-c', 'echo Hello Kubernetes!'], 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }, { 'secretRef': { 'name': 'secret_a' } }], 'resources': { 'requests': { 'memory': '1Gi', 'cpu': 1, 'ephemeral-storage': '2Gi' }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 1, 'ephemeral-storage': '4Gi' }, }, 'ports': [{ 'name': 'foo', 'containerPort': 1234 }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol' + str(self.static_uuid), 'readOnly': True }] }], 'volumes': [{ 'name': 'secretvol' + str(self.static_uuid), 'secret': { 'secretName': 'secret_b' } }], 'hostNetwork': False, 'imagePullSecrets': [{ 'name': 'pull_secret_a' }, { 'name': 'pull_secret_b' }], 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } } @mock.patch('uuid.uuid4') def test_gen_pod(self, mock_uuid): mock_uuid.return_value = self.static_uuid pod_generator = PodGenerator( labels={'app': 'myapp'}, name='myapp-pod', image_pull_secrets='pull_secret_a,pull_secret_b', image='busybox', envs=self.envs, cmds=['sh', '-c', 'echo Hello Kubernetes!'], security_context=k8s.V1PodSecurityContext( run_as_user=1000, fs_group=2000, ), namespace='default', ports=[k8s.V1ContainerPort(name='foo', container_port=1234)], configmaps=['configmap_a', 'configmap_b']) result = pod_generator.gen_pod() result = append_to_pod(result, self.secrets) result = self.resources.attach_to_pod(result) result_dict = self.k8s_client.sanitize_for_serialization(result) # sort result_dict['spec']['containers'][0]['env'].sort( key=lambda x: x['name']) result_dict['spec']['containers'][0]['envFrom'].sort( key=lambda x: list(x.values())[0]['name']) self.assertDictEqual(self.expected, result_dict) @mock.patch('uuid.uuid4') def test_gen_pod_extract_xcom(self, mock_uuid): mock_uuid.return_value = self.static_uuid pod_generator = PodGenerator( labels={'app': 'myapp'}, name='myapp-pod', image_pull_secrets='pull_secret_a,pull_secret_b', image='busybox', envs=self.envs, cmds=['sh', '-c', 'echo Hello Kubernetes!'], namespace='default', security_context=k8s.V1PodSecurityContext( run_as_user=1000, fs_group=2000, ), ports=[k8s.V1ContainerPort(name='foo', container_port=1234)], configmaps=['configmap_a', 'configmap_b'], extract_xcom=True) result = pod_generator.gen_pod() result = append_to_pod(result, self.secrets) result = self.resources.attach_to_pod(result) result_dict = self.k8s_client.sanitize_for_serialization(result) container_two = { 'name': 'airflow-xcom-sidecar', 'image': "alpine", 'command': ['sh', '-c', PodDefaults.XCOM_CMD], 'volumeMounts': [{ 'name': 'xcom', 'mountPath': '/airflow/xcom' }], 'resources': { 'requests': { 'cpu': '1m' } }, } self.expected['spec']['containers'].append(container_two) self.expected['spec']['containers'][0]['volumeMounts'].insert( 0, { 'name': 'xcom', 'mountPath': '/airflow/xcom' }) self.expected['spec']['volumes'].insert(0, { 'name': 'xcom', 'emptyDir': {} }) result_dict['spec']['containers'][0]['env'].sort( key=lambda x: x['name']) self.assertEqual(result_dict, self.expected) @mock.patch('uuid.uuid4') def test_from_obj(self, mock_uuid): mock_uuid.return_value = self.static_uuid result = PodGenerator.from_obj({ "KubernetesExecutor": { "annotations": { "test": "annotation" }, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": { "path": "/tmp/" }, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], "resources": { "requests": { "memory": "256Mi", "cpu": "500m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" }, "limits": { "memory": "512Mi", "cpu": "1000m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" } } } }) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': { 'test': 'annotation' }, }, 'spec': { 'containers': [{ 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume' }], "resources": { "requests": { "memory": "256Mi", "cpu": "500m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" }, "limits": { "memory": "512Mi", "cpu": "1000m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" } } }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result) @mock.patch('uuid.uuid4') def test_from_obj_with_resources_object(self, mock_uuid): mock_uuid.return_value = self.static_uuid result = PodGenerator.from_obj({ "KubernetesExecutor": { "annotations": { "test": "annotation" }, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": { "path": "/tmp/" }, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], "resources": { "requests": { "memory": "256Mi", "cpu": "500m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" }, "limits": { "memory": "512Mi", "cpu": "1000m", "ephemeral-storage": "2G", "nvidia.com/gpu": "0" } } } }) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': { 'test': 'annotation' }, }, 'spec': { 'containers': [{ 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume' }], 'resources': { 'limits': { 'cpu': '1000m', 'ephemeral-storage': '2G', 'memory': '512Mi', 'nvidia.com/gpu': '0' }, 'requests': { 'cpu': '500m', 'ephemeral-storage': '2G', 'memory': '256Mi', 'nvidia.com/gpu': '0' } }, }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result) @mock.patch('uuid.uuid4') def test_from_obj_with_resources(self, mock_uuid): self.maxDiff = None mock_uuid.return_value = self.static_uuid result = PodGenerator.from_obj({ "KubernetesExecutor": { "annotations": { "test": "annotation" }, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": { "path": "/tmp/" }, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], 'request_cpu': "200m", 'limit_cpu': "400m", 'request_memory': "500Mi", 'limit_memory': "1000Mi", 'limit_gpu': "2", 'request_ephemeral_storage': '2Gi', 'limit_ephemeral_storage': '4Gi', } }) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': { 'test': 'annotation' }, }, 'spec': { 'containers': [{ 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'limits': { 'cpu': '400m', 'ephemeral-storage': '4Gi', 'memory': '1000Mi', 'nvidia.com/gpu': "2", }, 'requests': { 'cpu': '200m', 'ephemeral-storage': '2Gi', 'memory': '500Mi', }, }, 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume' }], }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result) @mock.patch('uuid.uuid4') def test_from_obj_with_only_request_resources(self, mock_uuid): self.maxDiff = None mock_uuid.return_value = self.static_uuid result = PodGenerator.from_obj({ "KubernetesExecutor": { "annotations": { "test": "annotation" }, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": { "path": "/tmp/" }, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], 'request_cpu': "200m", 'request_memory': "500Mi", } }) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': { 'test': 'annotation' }, }, 'spec': { 'containers': [{ 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'requests': { 'cpu': '200m', 'memory': '500Mi', }, }, 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume' }], }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result) @mock.patch('uuid.uuid4') def test_reconcile_pods_empty_mutator_pod(self, mock_uuid): mock_uuid.return_value = self.static_uuid base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=[k8s.V1ContainerPort(name='port', container_port=2118)], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = None name = 'name1-' + self.static_uuid.hex base_pod.metadata.name = name result = PodGenerator.reconcile_pods(base_pod, mutator_pod) self.assertEqual(base_pod, result) mutator_pod = k8s.V1Pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) self.assertEqual(base_pod, result) @mock.patch('uuid.uuid4') def test_reconcile_pods(self, mock_uuid): mock_uuid.return_value = self.static_uuid base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=[k8s.V1ContainerPort(name='port', container_port=2118)], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = PodGenerator(envs={ 'key2': 'val2' }, image='', name='name2', cmds=['/bin/command2.sh', 'arg2'], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' }]).gen_pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'name2-' + self.static_uuid.hex }, 'spec': { 'containers': [{ 'args': [], 'command': ['/bin/command2.sh', 'arg2'], 'env': [{ 'name': 'key1', 'value': 'val1' }, { 'name': 'key2', 'value': 'val2' }], 'envFrom': [], 'image': 'image1', 'name': 'base', 'ports': [{ 'containerPort': 2118, 'name': 'port', }], 'volumeMounts': [{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }, { 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }] } }, result) @mock.patch('uuid.uuid4') def test_construct_pod_empty_worker_config(self, mock_uuid): mock_uuid.return_value = self.static_uuid executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '1m', 'memory': '1G' })) ])) worker_config = k8s.V1Pod() result = PodGenerator.construct_pod( 'dag_id', 'task_id', 'pod_id', 3, 'date', ['command'], executor_config, worker_config, 'namespace', 'uuid', ) sanitized_result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': self.metadata, 'spec': { 'containers': [{ 'args': [], 'command': ['command'], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'limits': { 'cpu': '1m', 'memory': '1G' } }, 'volumeMounts': [] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [] } }, sanitized_result) @mock.patch('uuid.uuid4') def test_construct_pod_empty_executor_config(self, mock_uuid): mock_uuid.return_value = self.static_uuid worker_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '1m', 'memory': '1G' })) ])) executor_config = None result = PodGenerator.construct_pod( 'dag_id', 'task_id', 'pod_id', 3, 'date', ['command'], executor_config, worker_config, 'namespace', 'uuid', ) sanitized_result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': self.metadata, 'spec': { 'containers': [{ 'args': [], 'command': ['command'], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'limits': { 'cpu': '1m', 'memory': '1G' } }, 'volumeMounts': [] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [] } }, sanitized_result) @mock.patch('uuid.uuid4') def test_construct_pod(self, mock_uuid): mock_uuid.return_value = self.static_uuid worker_config = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name='gets-overridden-by-dynamic-args', annotations={'should': 'stay'}), spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='doesnt-override', resources=k8s.V1ResourceRequirements( limits={ 'cpu': '1m', 'memory': '1G' }), security_context=k8s.V1SecurityContext( run_as_user=1)) ])) executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='doesnt-override-either', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '2m', 'memory': '2G' })) ])) result = PodGenerator.construct_pod( 'dag_id', 'task_id', 'pod_id', 3, 'date', ['command'], executor_config, worker_config, 'namespace', 'uuid', ) sanitized_result = self.k8s_client.sanitize_for_serialization(result) self.metadata.update({'annotations': {'should': 'stay'}}) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': self.metadata, 'spec': { 'containers': [{ 'args': [], 'command': ['command'], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'limits': { 'cpu': '2m', 'memory': '2G' } }, 'volumeMounts': [], 'securityContext': { 'runAsUser': 1 } }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [] } }, sanitized_result) @mock.patch('uuid.uuid4') def test_construct_pod_with_mutation(self, mock_uuid): mock_uuid.return_value = self.static_uuid worker_config = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name='gets-overridden-by-dynamic-args', annotations={'should': 'stay'}), spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='doesnt-override', resources=k8s.V1ResourceRequirements( limits={ 'cpu': '1m', 'memory': '1G' }), security_context=k8s.V1SecurityContext( run_as_user=1)) ])) executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='doesnt-override-either', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '2m', 'memory': '2G' })) ])) result = PodGenerator.construct_pod( 'dag_id', 'task_id', 'pod_id', 3, 'date', ['command'], executor_config, worker_config, 'namespace', 'uuid', ) sanitized_result = self.k8s_client.sanitize_for_serialization(result) self.metadata.update({'annotations': {'should': 'stay'}}) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': self.metadata, 'spec': { 'containers': [{ 'args': [], 'command': ['command'], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'resources': { 'limits': { 'cpu': '2m', 'memory': '2G' } }, 'volumeMounts': [], 'securityContext': { 'runAsUser': 1 } }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [] } }, sanitized_result) def test_merge_objects_empty(self): annotations = {'foo1': 'bar1'} base_obj = k8s.V1ObjectMeta(annotations=annotations) client_obj = None res = merge_objects(base_obj, client_obj) self.assertEqual(base_obj, res) client_obj = k8s.V1ObjectMeta() res = merge_objects(base_obj, client_obj) self.assertEqual(base_obj, res) client_obj = k8s.V1ObjectMeta(annotations=annotations) base_obj = None res = merge_objects(base_obj, client_obj) self.assertEqual(client_obj, res) base_obj = k8s.V1ObjectMeta() res = merge_objects(base_obj, client_obj) self.assertEqual(client_obj, res) def test_merge_objects(self): base_annotations = {'foo1': 'bar1'} base_labels = {'foo1': 'bar1'} client_annotations = {'foo2': 'bar2'} base_obj = k8s.V1ObjectMeta(annotations=base_annotations, labels=base_labels) client_obj = k8s.V1ObjectMeta(annotations=client_annotations) res = merge_objects(base_obj, client_obj) client_obj.labels = base_labels self.assertEqual(client_obj, res) def test_extend_object_field_empty(self): ports = [k8s.V1ContainerPort(container_port=1, name='port')] base_obj = k8s.V1Container(name='base_container', ports=ports) client_obj = k8s.V1Container(name='client_container') res = extend_object_field(base_obj, client_obj, 'ports') client_obj.ports = ports self.assertEqual(client_obj, res) base_obj = k8s.V1Container(name='base_container') client_obj = k8s.V1Container(name='base_container', ports=ports) res = extend_object_field(base_obj, client_obj, 'ports') self.assertEqual(client_obj, res) def test_extend_object_field_not_list(self): base_obj = k8s.V1Container(name='base_container', image='image') client_obj = k8s.V1Container(name='client_container') with self.assertRaises(ValueError): extend_object_field(base_obj, client_obj, 'image') base_obj = k8s.V1Container(name='base_container') client_obj = k8s.V1Container(name='client_container', image='image') with self.assertRaises(ValueError): extend_object_field(base_obj, client_obj, 'image') def test_extend_object_field(self): base_ports = [k8s.V1ContainerPort(container_port=1, name='base_port')] base_obj = k8s.V1Container(name='base_container', ports=base_ports) client_ports = [ k8s.V1ContainerPort(container_port=1, name='client_port') ] client_obj = k8s.V1Container(name='client_container', ports=client_ports) res = extend_object_field(base_obj, client_obj, 'ports') client_obj.ports = base_ports + client_ports self.assertEqual(client_obj, res) def test_reconcile_containers_empty(self): base_objs = [k8s.V1Container(name='base_container')] client_objs = [] res = PodGenerator.reconcile_containers(base_objs, client_objs) self.assertEqual(base_objs, res) client_objs = [k8s.V1Container(name='client_container')] base_objs = [] res = PodGenerator.reconcile_containers(base_objs, client_objs) self.assertEqual(client_objs, res) res = PodGenerator.reconcile_containers([], []) self.assertEqual(res, []) def test_reconcile_containers(self): base_ports = [k8s.V1ContainerPort(container_port=1, name='base_port')] base_objs = [ k8s.V1Container(name='base_container1', ports=base_ports), k8s.V1Container(name='base_container2', image='base_image'), ] client_ports = [ k8s.V1ContainerPort(container_port=2, name='client_port') ] client_objs = [ k8s.V1Container(name='client_container1', ports=client_ports), k8s.V1Container(name='client_container2', image='client_image'), ] res = PodGenerator.reconcile_containers(base_objs, client_objs) client_objs[0].ports = base_ports + client_ports self.assertEqual(client_objs, res) base_ports = [k8s.V1ContainerPort(container_port=1, name='base_port')] base_objs = [ k8s.V1Container(name='base_container1', ports=base_ports), k8s.V1Container(name='base_container2', image='base_image'), ] client_ports = [ k8s.V1ContainerPort(container_port=2, name='client_port') ] client_objs = [ k8s.V1Container(name='client_container1', ports=client_ports), k8s.V1Container(name='client_container2', stdin=True), ] res = PodGenerator.reconcile_containers(base_objs, client_objs) client_objs[0].ports = base_ports + client_ports client_objs[1].image = 'base_image' self.assertEqual(client_objs, res) def test_reconcile_specs_empty(self): base_spec = k8s.V1PodSpec(containers=[]) client_spec = None res = PodGenerator.reconcile_specs(base_spec, client_spec) self.assertEqual(base_spec, res) base_spec = None client_spec = k8s.V1PodSpec(containers=[]) res = PodGenerator.reconcile_specs(base_spec, client_spec) self.assertEqual(client_spec, res) def test_reconcile_specs(self): base_objs = [ k8s.V1Container(name='base_container1', image='base_image') ] client_objs = [k8s.V1Container(name='client_container1')] base_spec = k8s.V1PodSpec(priority=1, active_deadline_seconds=100, containers=base_objs) client_spec = k8s.V1PodSpec(priority=2, hostname='local', containers=client_objs) res = PodGenerator.reconcile_specs(base_spec, client_spec) client_spec.containers = [ k8s.V1Container(name='client_container1', image='base_image') ] client_spec.active_deadline_seconds = 100 self.assertEqual(client_spec, res) def test_deserialize_model_file(self): fixture = 'tests/kubernetes/pod.yaml' result = PodGenerator.deserialize_model_file(fixture) sanitized_res = self.k8s_client.sanitize_for_serialization(result) self.assertEqual(sanitized_res, self.deserialize_result) def test_deserialize_model_string(self): fixture = """ apiVersion: v1 kind: Pod metadata: name: memory-demo namespace: mem-example spec: containers: - name: memory-demo-ctr image: polinux/stress resources: limits: memory: "200Mi" requests: memory: "100Mi" command: ["stress"] args: ["--vm", "1", "--vm-bytes", "150M", "--vm-hang", "1"] """ result = PodGenerator.deserialize_model_file(fixture) sanitized_res = self.k8s_client.sanitize_for_serialization(result) self.assertEqual(sanitized_res, self.deserialize_result) def test_validate_pod_generator(self): with self.assertRaises(AirflowConfigException): PodGenerator(image='k', pod=k8s.V1Pod()) with self.assertRaises(AirflowConfigException): PodGenerator(pod=k8s.V1Pod(), pod_template_file='k') with self.assertRaises(AirflowConfigException): PodGenerator(image='k', pod_template_file='k') PodGenerator(image='k') PodGenerator(pod_template_file='tests/kubernetes/pod.yaml') PodGenerator(pod=k8s.V1Pod()) def test_add_custom_label(self): from kubernetes.client import models as k8s pod = PodGenerator.construct_pod( namespace="test", worker_uuid="test", pod_id="test", dag_id="test", task_id="test", try_number=1, date="23-07-2020", command="test", kube_executor_config=None, worker_config=k8s.V1Pod(metadata=k8s.V1ObjectMeta( labels={"airflow-test": "airflow-task-pod"}, annotations={"my.annotation": "foo"}))) self.assertIn("airflow-test", pod.metadata.labels) self.assertIn("my.annotation", pod.metadata.annotations)