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
assert 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'
assert client_objs == res
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
assert client_obj == res
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 test_port(self):
port = k8s.V1ContainerPort(
name='http',
container_port=80,
)
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
ports=[port],
)
context = create_context(k)
k.execute(context=context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['ports'] = [{
'name':
'http',
'containerPort':
80
}]
self.assertEqual(self.expected_pod, actual_pod)
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)
def to_k8s_client_obj(self):
"""
Converts to k8s object.
:rtype: object
"""
return k8s.V1ContainerPort(name=self.name,
container_port=self.container_port)
def test_port_to_k8s_client_obj(self):
port = Port('http', 80)
self.assertEqual(
port.to_k8s_client_obj(),
k8s.V1ContainerPort(
name='http',
container_port=80
)
)
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_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_convert_to_airflow_pod(self):
input_pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name="foo", namespace="bar"),
spec=k8s.V1PodSpec(
init_containers=[
k8s.V1Container(name="init-container",
volume_mounts=[
k8s.V1VolumeMount(mount_path="/tmp",
name="init-secret")
])
],
containers=[
k8s.V1Container(
name="base",
command=["foo"],
image="myimage",
env=[
k8s.V1EnvVar(
name="AIRFLOW_SECRET",
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
name="ai", key="secret_key")))
],
ports=[
k8s.V1ContainerPort(
name="myport",
container_port=8080,
)
],
volume_mounts=[
k8s.V1VolumeMount(name="myvolume",
mount_path="/tmp/mount",
read_only="True"),
k8s.V1VolumeMount(name='airflow-config',
mount_path='/config',
sub_path='airflow.cfg',
read_only=True),
k8s.V1VolumeMount(name="airflow-secret",
mount_path="/opt/mount",
read_only=True)
])
],
security_context=k8s.V1PodSecurityContext(
run_as_user=0,
fs_group=0,
),
volumes=[
k8s.V1Volume(name="myvolume"),
k8s.V1Volume(
name="airflow-config",
config_map=k8s.V1ConfigMap(data="airflow-data")),
k8s.V1Volume(name="airflow-secret",
secret=k8s.V1SecretVolumeSource(
secret_name="secret-name", )),
k8s.V1Volume(name="init-secret",
secret=k8s.V1SecretVolumeSource(
secret_name="init-secret", ))
]))
result_pod = _convert_to_airflow_pod(input_pod)
expected = Pod(
name="foo",
namespace="bar",
envs={},
init_containers=[{
'name':
'init-container',
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
cmds=["foo"],
image="myimage",
ports=[Port(name="myport", container_port=8080)],
volume_mounts=[
VolumeMount(name="myvolume",
mount_path="/tmp/mount",
sub_path=None,
read_only="True"),
VolumeMount(name="airflow-config",
read_only=True,
mount_path="/config",
sub_path="airflow.cfg"),
VolumeMount(name="airflow-secret",
mount_path="/opt/mount",
sub_path=None,
read_only=True)
],
secrets=[Secret("env", "AIRFLOW_SECRET", "ai", "secret_key")],
security_context={
'fsGroup': 0,
'runAsUser': 0
},
volumes=[
Volume(name="myvolume", configs={'name': 'myvolume'}),
Volume(name="airflow-config",
configs={
'configMap': {
'data': 'airflow-data'
},
'name': 'airflow-config'
}),
Volume(name='airflow-secret',
configs={
'name': 'airflow-secret',
'secret': {
'secretName': 'secret-name'
}
}),
Volume(name='init-secret',
configs={
'name': 'init-secret',
'secret': {
'secretName': 'init-secret'
}
})
],
)
expected_dict = expected.as_dict()
result_dict = result_pod.as_dict()
print(result_pod.volume_mounts)
parsed_configs = self.pull_out_volumes(result_dict)
result_dict['volumes'] = parsed_configs
self.assertEqual(result_dict['secrets'], expected_dict['secrets'])
self.assertDictEqual(expected_dict, result_dict)
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': 'apache/airflow:stress-2020.07.10-1.0.4',
'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.execution_date = parser.parse('2020-08-24 00:00:00.000000')
self.execution_date_label = datetime_to_label_safe_datestring(self.execution_date)
self.dag_id = 'dag_id'
self.task_id = 'task_id'
self.try_number = 3
self.labels = {
'airflow-worker': 'uuid',
'dag_id': self.dag_id,
'execution_date': self.execution_date_label,
'task_id': self.task_id,
'try_number': str(self.try_number),
'airflow_version': __version__.replace('+', '-'),
'kubernetes_executor': 'True',
}
self.annotations = {
'dag_id': self.dag_id,
'task_id': self.task_id,
'execution_date': self.execution_date.isoformat(),
'try_number': str(self.try_number),
}
self.metadata = {
'labels': self.labels,
'name': 'pod_id-' + self.static_uuid.hex,
'namespace': 'namespace',
'annotations': self.annotations,
}
self.resources = k8s.V1ResourceRequirements(
requests={
"cpu": 1,
"memory": "1Gi",
"ephemeral-storage": "2Gi",
},
limits={"cpu": 2, "memory": "2Gi", "ephemeral-storage": "4Gi", 'nvidia.com/gpu': 1},
)
self.k8s_client = ApiClient()
self.expected = k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
namespace="default",
name='myapp-pod-' + self.static_uuid.hex,
labels={'app': 'myapp'},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name='base',
image='busybox',
command=['sh', '-c', 'echo Hello Kubernetes!'],
env=[
k8s.V1EnvVar(name='ENVIRONMENT', value='prod'),
k8s.V1EnvVar(
name="LOG_LEVEL",
value='warning',
),
k8s.V1EnvVar(
name='TARGET',
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(name='secret_b', key='source_b')
),
),
],
env_from=[
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_a')),
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_b')),
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name='secret_a')),
],
ports=[k8s.V1ContainerPort(name="foo", container_port=1234)],
resources=k8s.V1ResourceRequirements(
requests={'memory': '100Mi'},
limits={
'memory': '200Mi',
},
),
)
],
security_context=k8s.V1PodSecurityContext(
fs_group=2000,
run_as_user=1000,
),
host_network=True,
image_pull_secrets=[
k8s.V1LocalObjectReference(name="pull_secret_a"),
k8s.V1LocalObjectReference(name="pull_secret_b"),
],
),
)
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'
}],
labels={
"foo": "bar"
},
volume_mounts=[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}],
).gen_pod()
mutator_pod = PodGenerator(envs={
'key2': 'val2'
},
image='',
name='name2',
labels={
"bar": "baz"
},
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,
'labels': {
'foo': 'bar',
"bar": "baz"
}
},
'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-volume1'
}, {
'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)
secret_all_keys = Secret('env', None, 'airflow-secrets-2')
volume_mount = k8s.V1VolumeMount(
name='test-volume', mount_path='/root/mount_file', sub_path=None, read_only=True
)
configmaps = [
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='test-configmap-1')),
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='test-configmap-2')),
]
volume = k8s.V1Volume(
name='test-volume',
persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name='test-volume'),
)
port = k8s.V1ContainerPort(name='http', container_port=80)
init_container_volume_mounts = [
k8s.V1VolumeMount(mount_path='/etc/foo', name='test-volume', sub_path=None, read_only=True)
]
init_environments = [k8s.V1EnvVar(name='key1', value='value1'), k8s.V1EnvVar(name='key2', value='value2')]
init_container = k8s.V1Container(
name="init-container",
image="ubuntu:16.04",
env=init_environments,
volume_mounts=init_container_volume_mounts,
command=["bash", "-cx"],
args=["echo 10"],
)
def to_k8s_client_obj(self) -> k8s.V1ContainerPort:
"""Converts to k8s client object"""
return k8s.V1ContainerPort(name=self.name,
container_port=self.container_port)
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'
}]
}
})
read_only=True)
configmaps = [
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name="test-configmap-1")),
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name="test-configmap-2")),
]
volume = k8s.V1Volume(
name="test-volume",
persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(
claim_name="test-volume"),
)
port = k8s.V1ContainerPort(name="http", container_port=80)
init_container_volume_mounts = [
k8s.V1VolumeMount(mount_path="/etc/foo",
name="test-volume",
sub_path=None,
read_only=True)
]
init_environments = [
k8s.V1EnvVar(name="key1", value="value1"),
k8s.V1EnvVar(name="key2", value="value2"),
]
init_container = k8s.V1Container(
name="init-container",
def to_k8s_client_obj(self) -> k8s.V1ContainerPort:
return k8s.V1ContainerPort(name=self.name,
container_port=self.container_port)
