def test_gen_pod_extract_xcom(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
path = sys.path[0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml'
pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True)
result = pod_generator.gen_pod()
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)
base_container: k8s.V1Container = self.expected.spec.containers[0]
base_container.volume_mounts = base_container.volume_mounts or []
base_container.volume_mounts.append(k8s.V1VolumeMount(name="xcom", mount_path="/airflow/xcom"))
self.expected.spec.containers[0] = base_container
self.expected.spec.volumes = self.expected.spec.volumes or []
self.expected.spec.volumes.append(
k8s.V1Volume(
name='xcom',
empty_dir={},
)
)
result_dict = self.k8s_client.sanitize_for_serialization(result)
expected_dict = self.k8s_client.sanitize_for_serialization(self.expected)
assert result_dict == expected_dict
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)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, try_number, airflow_command) -> k8s.V1Pod:
"""Creates POD."""
pod_generator = PodGenerator(
namespace=namespace,
name=pod_id,
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy,
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date,
'try_number': str(try_number),
'airflow_version': airflow_version.replace('+', '-'),
'kubernetes_executor': 'True',
},
cmds=airflow_command,
volumes=self._get_volumes(),
volume_mounts=self._get_volume_mounts(),
init_containers=self._get_init_containers(),
annotations=self.kube_config.kube_annotations,
affinity=self.kube_config.kube_affinity,
tolerations=self.kube_config.kube_tolerations,
envs=self._get_environment(),
node_selectors=self.kube_config.kube_node_selectors,
service_account_name=self.kube_config.worker_service_account_name,
)
pod = pod_generator.gen_pod()
pod.spec.containers[0].env_from = pod.spec.containers[0].env_from or []
pod.spec.containers[0].env_from.extend(self._get_env_from())
pod.spec.security_context = self._get_security_context()
return append_to_pod(pod, self._get_secrets())
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)
def test_reconcile_pods_empty_mutator_pod(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
path = sys.path[0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml'
pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True)
base_pod = pod_generator.gen_pod()
mutator_pod = None
name = 'name1-' + self.static_uuid.hex
base_pod.metadata.name = name
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
assert base_pod == result
mutator_pod = k8s.V1Pod()
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
assert base_pod == result
def as_pod(self) -> k8s.V1Pod:
"""Creates POD."""
pod_generator = PodGenerator(
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy,
image_pull_secrets=self.kube_config.image_pull_secrets,
volumes=self._get_volumes(),
volume_mounts=self._get_volume_mounts(),
init_containers=self._get_init_containers(),
annotations=self.kube_config.kube_annotations,
affinity=self.kube_config.kube_affinity,
tolerations=self.kube_config.kube_tolerations,
envs=self._get_environment(),
node_selectors=self.kube_config.kube_node_selectors,
service_account_name=self.kube_config.worker_service_account_name,
)
pod = pod_generator.gen_pod()
pod.spec.containers[0].env_from = pod.spec.containers[0].env_from or []
pod.spec.containers[0].env_from.extend(self._get_env_from())
pod.spec.security_context = self._get_security_context()
return append_to_pod(pod, self._get_secrets())
