def test_extract_labels(self): # Test when labels are not empty labels = {'label_a': 'val_a', 'label_b': 'val_b'} pod = Pod('v3.14', {}, [], labels=labels) self.expected['metadata']['labels'] = labels KubernetesRequestFactory.extract_labels(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_attach_volume_mounts(self): # Test when volumes is not empty list volume_mounts = ['vol_a', 'vol_b'] pod = Pod('v3.14', {}, [], volume_mounts=volume_mounts) self.expected['spec']['containers'][0]['volumeMounts'] = volume_mounts KubernetesRequestFactory.attach_volume_mounts(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_volume_secrets(self): # Test when secrets is not empty secrets = [ Secret('volume', 'KEY1', 's1', 'key-1'), Secret('env', 'KEY2', 's2'), Secret('volume', 'KEY3', 's3', 'key-2') ] pod = Pod('v3.14', {}, [], secrets=secrets) self.expected['spec']['containers'][0]['volumeMounts'] = [{ 'mountPath': 'KEY1', 'name': 'secretvol0', 'readOnly': True }, { 'mountPath': 'KEY3', 'name': 'secretvol1', 'readOnly': True }] self.expected['spec']['volumes'] = [{ 'name': 'secretvol0', 'secret': { 'secretName': 's1' } }, { 'name': 'secretvol1', 'secret': { 'secretName': 's3' } }] KubernetesRequestFactory.extract_volume_secrets(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_annotations(self): # Test when annotations are not empty annotations = {'annot_a': 'val_a', 'annot_b': 'val_b'} pod = Pod('v3.14', {}, [], annotations=annotations) self.expected['metadata']['annotations'] = annotations KubernetesRequestFactory.extract_annotations(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_affinity(self): # Test when affinity is not empty affinity = {'podAffinity': 'requiredDuringSchedulingIgnoredDuringExecution'} pod = Pod('v3.14', {}, [], affinity=affinity) self.expected['spec']['affinity'] = affinity KubernetesRequestFactory.extract_affinity(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_node_selector(self): # Test when affinity is not empty node_selectors = {'disktype': 'ssd', 'accelerator': 'nvidia-tesla-p100'} pod = Pod('v3.14', {}, [], node_selectors=node_selectors) self.expected['spec']['nodeSelector'] = node_selectors KubernetesRequestFactory.extract_node_selector(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_service_account_name(self): service_account_name = 'service_account_name' pod = Pod('v3.14', {}, [], service_account_name=service_account_name) self.expected['spec']['serviceAccountName'] = service_account_name KubernetesRequestFactory.extract_service_account_name( pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_image_pull_policy(self): # Test when pull policy is not none pull_policy = 'IfNotPresent' pod = Pod('v3.14', {}, [], image_pull_policy=pull_policy) KubernetesRequestFactory.extract_image_pull_policy(pod, self.input_req) self.expected['spec']['containers'][0]['imagePullPolicy'] = pull_policy self.assertEqual(self.input_req, self.expected)
def test_security_context(self): security_context = { 'runAsUser': 1000, 'fsGroup': 2000 } pod = Pod('v3.14', {}, [], security_context=security_context) self.expected['spec']['securityContext'] = security_context KubernetesRequestFactory.extract_security_context(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id, execution_date, try_number, airflow_command, kube_executor_config): volumes_dict, volume_mounts_dict = self._get_volumes_and_mounts() worker_init_container_spec = self._get_init_containers() resources = Resources( request_memory=kube_executor_config.request_memory, request_cpu=kube_executor_config.request_cpu, limit_memory=kube_executor_config.limit_memory, limit_cpu=kube_executor_config.limit_cpu, limit_gpu=kube_executor_config.limit_gpu) gcp_sa_key = kube_executor_config.gcp_service_account_key annotations = dict(kube_executor_config.annotations ) or self.kube_config.kube_annotations if gcp_sa_key: annotations['iam.cloud.google.com/service-account'] = gcp_sa_key volumes = [value for value in volumes_dict.values() ] + kube_executor_config.volumes volume_mounts = [value for value in volume_mounts_dict.values() ] + kube_executor_config.volume_mounts affinity = kube_executor_config.affinity or self.kube_config.kube_affinity tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations return Pod( namespace=namespace, name=pod_id, image=kube_executor_config.image or self.kube_config.kube_image, image_pull_policy=(kube_executor_config.image_pull_policy or self.kube_config.kube_image_pull_policy), cmds=airflow_command, labels=self._get_labels( kube_executor_config.labels, { 'airflow-worker': worker_uuid, 'dag_id': dag_id, 'task_id': task_id, 'execution_date': execution_date, 'try_number': str(try_number), }), envs=self._get_environment(), secrets=self._get_secrets(), service_account_name=self.kube_config.worker_service_account_name, image_pull_secrets=self.kube_config.image_pull_secrets, init_containers=worker_init_container_spec, volumes=volumes, volume_mounts=volume_mounts, resources=resources, annotations=annotations, node_selectors=(kube_executor_config.node_selectors or self.kube_config.kube_node_selectors), affinity=affinity, tolerations=tolerations, security_context=self._get_security_context(), configmaps=self._get_configmaps(), host_aliases=self._get_host_aliases())
def test_extract_image_pull_secrets(self): image_pull_secrets = 'secret_a,secret_b,secret_c' pod = Pod('v3.14', {}, [], image_pull_secrets=image_pull_secrets) self.expected['spec']['imagePullSecrets'] = [ {'name': 'secret_a'}, {'name': 'secret_b'}, {'name': 'secret_c'}, ] KubernetesRequestFactory.extract_image_pull_secrets(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_tolerations(self): tolerations = [{ 'key': 'key', 'operator': 'Equal', 'value': 'value', 'effect': 'NoSchedule' }] pod = Pod('v3.14', {}, [], tolerations=tolerations) self.expected['spec']['tolerations'] = tolerations KubernetesRequestFactory.extract_tolerations(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_env_and_secrets(self): # Test when secrets and envs are not empty secrets = [ Secret('env', None, 's1'), Secret('volume', 'KEY2', 's2', 'key-2'), Secret('env', None, 's3') ] envs = { 'ENV1': 'val1', 'ENV2': 'val2' } configmaps = ['configmap_a', 'configmap_b'] pod_runtime_envs = [PodRuntimeInfoEnv("ENV3", "status.podIP")] pod = Pod( image='v3.14', envs=envs, cmds=[], secrets=secrets, configmaps=configmaps, pod_runtime_info_envs=pod_runtime_envs) self.expected['spec']['containers'][0]['env'] = [ {'name': 'ENV1', 'value': 'val1'}, {'name': 'ENV2', 'value': 'val2'}, { 'name': 'ENV3', 'valueFrom': { 'fieldRef': { 'fieldPath': 'status.podIP' } } } ] self.expected['spec']['containers'][0]['envFrom'] = [{ 'secretRef': { 'name': 's1' } }, { 'secretRef': { 'name': 's3' } }, { 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }] KubernetesRequestFactory.extract_env_and_secrets(pod, self.input_req) self.input_req['spec']['containers'][0]['env'].sort(key=lambda x: x['name']) self.assertEqual(self.input_req, self.expected)
def test_extract_limits_resources(self): # Test when resources is not empty resources = Resources(limit_memory='2Gi', limit_cpu=2) pod = Pod('v3.14', {}, [], resources=resources) self.expected['spec']['containers'][0]['resources'] = { 'limits': { 'memory': '2Gi', 'cpu': 2 } } KubernetesRequestFactory.extract_resources(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_pod_mutation_hook(self): """ Tests that pods are mutated by the pod_mutation_hook function in airflow_local_settings. """ with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"): from airflow import settings settings.import_local_settings() # pylint: ignore pod = Pod(image="ubuntu", envs={}, cmds=['echo "1"']) settings.pod_mutation_hook(pod) assert pod.namespace == 'airflow-tests'
def make_pod(self, namespace, image, pod_id, cmds, arguments, labels): volumes, volume_mounts = self._get_volumes_and_mounts() worker_init_container_spec = self._get_init_containers() return Pod( namespace=namespace, name=pod_id + "-" + str(uuid.uuid4())[:8], image=image, cmds=cmds, args=arguments, labels=labels, envs={}, secrets=[], # service_account_name=self.kube_config.worker_service_account_name, # image_pull_secrets=self.kube_config.image_pull_secrets, init_containers=worker_init_container_spec, volumes=volumes, volume_mounts=volume_mounts, resources=None)
def test_extract_env_and_secrets(self): # Test when secrets and envs are not empty secrets = [ Secret('env', None, 's1'), Secret('volume', 'KEY2', 's2', 'key-2'), Secret('env', None, 's3') ] envs = {'ENV1': 'val1', 'ENV2': 'val2'} configmaps = ['configmap_a', 'configmap_b'] pod = Pod('v3.14', envs, [], secrets=secrets, configmaps=configmaps) self.expected['spec']['containers'][0]['env'] = [ { 'name': 'ENV1', 'value': 'val1' }, { 'name': 'ENV2', 'value': 'val2' }, ] self.expected['spec']['containers'][0]['envFrom'] = [{ 'secretRef': { 'name': 's1' } }, { 'secretRef': { 'name': 's3' } }, { 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }] KubernetesRequestFactory.extract_env_and_secrets(pod, self.input_req) self.input_req['spec']['containers'][0]['env'].sort( key=lambda x: x['name']) self.assertEqual(self.input_req, self.expected)
def test_extract_all_resources(self): # Test when resources is not empty resources = Resources(request_memory='1Gi', request_cpu=1, limit_memory='2Gi', limit_cpu=2, limit_gpu=3) pod = Pod('v3.14', {}, [], resources=resources) self.expected['spec']['containers'][0]['resources'] = { 'requests': { 'memory': '1Gi', 'cpu': 1 }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 3 } } KubernetesRequestFactory.extract_resources(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_name(self): name = 'pod-name' pod = Pod('v3.14', {}, [], name=name) self.expected['metadata']['name'] = name KubernetesRequestFactory.extract_name(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def setUp(self): self.simple_pod_request_factory = SimplePodRequestFactory() self.xcom_pod_request_factory = ExtractXcomPodRequestFactory() self.pod = Pod( image='busybox', envs={ 'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning' }, name='myapp-pod', cmds=['sh', '-c', 'echo Hello Kubernetes!'], labels={'app': 'myapp'}, image_pull_secrets='pull_secret_a,pull_secret_b', configmaps=['configmap_a', 'configmap_b'], ports=[{ 'name': 'foo', 'containerPort': 1234 }], 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'), ], security_context={ 'runAsUser': 1000, 'fsGroup': 2000, }) self.maxDiff = None self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod', 'labels': { 'app': 'myapp' }, 'annotations': {} }, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'command': ['sh', '-c', 'echo Hello Kubernetes!'], 'imagePullPolicy': 'IfNotPresent', 'args': [], 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'secretRef': { 'name': 'secret_a' } }, { 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }], 'ports': [{ 'name': 'foo', 'containerPort': 1234 }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol0', 'readOnly': True }] }], 'restartPolicy': 'Never', 'nodeSelector': {}, 'volumes': [{ 'name': 'secretvol0', 'secret': { 'secretName': 'secret_b' } }], 'imagePullSecrets': [{ 'name': 'pull_secret_a' }, { 'name': 'pull_secret_b' }], 'affinity': {}, 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } }
def test_extract_cmds(self): cmds = ['test-cmd.sh'] pod = Pod('v3.14', {}, cmds) KubernetesRequestFactory.extract_cmds(pod, self.input_req) self.expected['spec']['containers'][0]['command'] = cmds self.assertEqual(self.input_req, self.expected)
def test_extract_args(self): args = ['test_arg.sh'] pod = Pod('v3.14', {}, [], args=args) KubernetesRequestFactory.extract_args(pod, self.input_req) self.expected['spec']['containers'][0]['args'] = args self.assertEqual(self.input_req, self.expected)
def test_extract_image(self): image = 'v3.14' pod = Pod(image, {}, []) KubernetesRequestFactory.extract_image(pod, self.input_req) self.expected['spec']['containers'][0]['image'] = image self.assertEqual(self.input_req, self.expected)
def test_identity(self, name): kube_request_factory_func = getattr(KubernetesRequestFactory, name) pod = Pod('v3.14', {}, []) kube_request_factory_func(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_hostnetwork(self): hostnetwork = True pod = Pod('v3.14', {}, [], hostnetwork=hostnetwork) self.expected['spec']['hostNetwork'] = hostnetwork KubernetesRequestFactory.extract_hostnetwork(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def test_extract_init_containers(self): init_container = 'init_container' pod = Pod('v3.14', {}, [], init_containers=init_container) self.expected['spec']['initContainers'] = init_container KubernetesRequestFactory.extract_init_containers(pod, self.input_req) self.assertEqual(self.input_req, self.expected)