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_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_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_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_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_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_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_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_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_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_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_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_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 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 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_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_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(self):
# Test when resources is not empty
resources = Resources(limit_memory='1Gi', limit_cpu=1)
pod = Pod('v3.14', {}, [], resources=resources)
self.expected['spec']['containers'][0]['resources'] = {
'limits': {
'memory': '1Gi',
'cpu': 1
}
}
KubernetesRequestFactory.extract_resources(pod, self.input_req)
self.assertEqual(self.expected, self.input_req)
def test_add_secret_to_env(self):
secret = Secret('env', 'target', 'my-secret', 'KEY')
secret_list = []
self.expected = [{
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'name': 'my-secret',
'key': 'KEY'
}
}
}]
KubernetesRequestFactory.add_secret_to_env(secret_list, secret)
self.assertListEqual(secret_list, self.expected)
def test_add_secret_to_env(self):
secret = Secret('env', 'target', 'my-secret', 'KEY')
secret_list = []
self.expected = [{
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'name': 'my-secret',
'key': 'KEY'
}
}
}]
KubernetesRequestFactory.add_secret_to_env(secret_list, secret)
self.assertListEqual(secret_list, self.expected)
def test_extract_resources(self):
# Test when resources is not empty
resources = Resources('1Gi', 1, '2Gi', 2)
pod = Pod('v3.14', {}, [], resources=resources)
self.expected['spec']['containers'][0]['resources'] = {
'requests': {
'memory': '1Gi',
'cpu': 1
},
'limits': {
'memory': '2Gi',
'cpu': 2
},
}
KubernetesRequestFactory.extract_resources(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def extract_env_and_secrets(pod, req):
KubernetesRequestFactory.extract_env_and_secrets(pod, req)
env = req['spec']['containers'][0]['env'].copy()
env = [i for i in env if not i["name"].startswith("AIRFLOW")]
kube_client = get_kube_client()
if configuration.conf.getboolean("kubernetes", "in_cluster"):
pod_config = kube_client.read_namespaced_pod(
name=os.getenv("HOSTNAME"),
namespace=configuration.conf.get("kubernetes", "namespace"))
else:
pods = kube_client.list_namespaced_pod(
namespace=configuration.conf.get("kubernetes", "namespace"),
label_selector="component=airflow-scheduler")
pod_config = pods.items[0]
self_env = [
convert_dict_key_case(e.to_dict())
for e in pod_config.spec.containers[0].env
if e.name not in ["AIRFLOW_COMMAND", "AIRFLOW_EXECUTOR"]
]
self_env.append({"name": "AIRFLOW_EXECUTOR", "value": "LocalExecutor"})
if configuration.conf.has_option("core", "worker_logging_level"):
self_env.append({
"name":
"AIRFLOW__CORE__LOGGING_LEVEL",
"value":
configuration.conf.get("core", "worker_logging_level")
})
self_env.append({
"name":
"AIRFLOW__CORE__FAB_LOGGING_LEVEL",
"value":
configuration.conf.get("core", "worker_logging_level")
})
req['spec']['containers'][0]['env'] = self_env + env
self_env_from = [
convert_dict_key_case(e.to_dict())
for e in pod_config.spec.containers[0].env_from
]
req['spec']['containers'][0]['envFrom'] = self_env_from
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_env_and_secret_order(self):
envs = {
'ENV': 'val1',
}
pod_runtime_envs = [PodRuntimeInfoEnv('RUNTIME_ENV', 'status.podIP')]
pod = Pod(image='v3.14',
envs=envs,
cmds=[],
pod_runtime_info_envs=pod_runtime_envs)
self.expected['spec']['containers'][0]['env'] = [{
'name': 'RUNTIME_ENV',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
}, {
'name': 'ENV',
'value': 'val1'
}]
KubernetesRequestFactory.extract_env_and_secrets(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 = 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 extract_env_and_secrets(pod, req):
"""Overrides KubernetesRequestFactory to support variable references of secrets
see https://github.com/apache/airflow/issues/8275"""
envs_from_key_secrets = [
env for env in pod.secrets if env.deploy_type == 'env' and env.key is not None
]
if len(pod.envs) > 0 or len(envs_from_key_secrets) > 0 or len(pod.pod_runtime_info_envs) > 0:
env = []
for runtime_info in pod.pod_runtime_info_envs:
KubernetesRequestFactory.add_runtime_info_env(env, runtime_info)
for secret in envs_from_key_secrets:
KubernetesRequestFactory.add_secret_to_env(env, secret)
for k in pod.envs.keys():
env.append({'name': k, 'value': pod.envs[k]})
req['spec']['containers'][0]['env'] = env
KubernetesRequestFactory._apply_env_from(pod, req)
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)
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 extract_labels(pod, req):
KubernetesRequestFactory.extract_labels(pod, req)
req['metadata']['labels']['cluster'] = 'airflow'
req['metadata']['labels']['component'] = 'airflow-worker'
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_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 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_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)
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_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 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_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_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_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)