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)
