def test_image_pull_secrets_correctly_set(self, mock_client, monitor_mock,
start_mock):
from airflow.utils.state import State
fake_pull_secrets = "fakeSecret"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
image_pull_secrets=fake_pull_secrets,
cluster_context='default',
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
self.assertEqual(start_mock.call_args[0][0].spec.image_pull_secrets,
[k8s.V1LocalObjectReference(name=fake_pull_secrets)])
def test_pod_priority_class_name(self, mock_client, monitor_mock,
start_mock): # pylint: disable=unused-argument
"""Test ability to assign priorityClassName to pod"""
priority_class_name = "medium-test"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
priority_class_name=priority_class_name,
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['priorityClassName'] = priority_class_name
assert self.expected_pod == actual_pod
def test_randomize_pod_name(self, mock_client, monitor_mock, start_mock):
from airflow.utils.state import State
name_base = 'test'
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
assert start_mock.call_args[0][0].metadata.name.startswith(name_base)
assert start_mock.call_args[0][0].metadata.name != name_base
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}]
assert self.expected_pod == actual_pod
def test_pod_failure(self):
"""
Tests that the task fails when a pod reports a failure
"""
bad_internal_command = ["foobar 10 "]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=bad_internal_command,
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,
)
with pytest.raises(AirflowException):
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0][
'args'] = bad_internal_command
assert self.expected_pod == actual_pod
def test_faulty_service_account(self):
"""pod creation should fail when service account does not exist"""
service_account = "foobar"
namespace = "default"
k = KubernetesPodOperator(
namespace=namespace,
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
startup_timeout_seconds=5,
service_account_name=service_account,
)
context = create_context(k)
pod = k.build_pod_request_obj(context)
with pytest.raises(
ApiException, match=f"error looking up service account {namespace}/{service_account}"
):
k.get_or_create_pod(pod, context)
def test_fs_group(self):
security_context = {
'securityContext': {
'fsGroup': 1000,
}
}
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
security_context=security_context,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['securityContext'] = security_context
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_name(self):
pod_name_too_long = "a" * 221
with pytest.raises(AirflowException):
KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=pod_name_too_long,
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
)
def test_image_pull_secrets_correctly_set(self, mock_client, await_pod_completion_mock, create_mock):
fake_pull_secrets = "fakeSecret"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
image_pull_secrets=fake_pull_secrets,
cluster_context='default',
)
mock_pod = MagicMock()
mock_pod.status.phase = 'Succeeded'
await_pod_completion_mock.return_value = mock_pod
context = create_context(k)
k.execute(context=context)
assert create_mock.call_args[1]['pod'].spec.image_pull_secrets == [
k8s.V1LocalObjectReference(name=fake_pull_secrets)
]
def test_image_pull_policy_not_set(self, mock_client, monitor_mock,
start_mock):
from airflow.utils.state import State
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
self.assertEqual(
start_mock.call_args[0][0].spec.containers[0].image_pull_policy,
'IfNotPresent',
)
def test_pod_dnspolicy(self):
dns_policy = "ClusterFirstWithHostNet"
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,
hostnetwork=True,
dnspolicy=dns_policy,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['hostNetwork'] = True
self.expected_pod['spec']['dnsPolicy'] = dns_policy
assert self.expected_pod['spec'] == actual_pod['spec']
assert self.expected_pod['metadata']['labels'] == actual_pod[
'metadata']['labels']
def test_port(self):
port = Port('http', 80)
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
ports=[port],
)
context = self.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_run_as_user_root(self):
security_context = {
'securityContext': {
'runAsUser': 0,
}
}
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,
security_context=security_context,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['securityContext'] = security_context
assert self.expected_pod == actual_pod
def test_no_need_to_describe_pod_on_success(self, mock_client, monitor_mock, start_mock):
from airflow.utils.state import State
name_base = 'test'
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
assert not mock_client.return_value.read_namespaced_pod.called
def get_op():
return KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["exit 1"],
labels={"foo": "bar"},
name="test",
task_id=name,
in_cluster=False,
do_xcom_push=False,
is_delete_operator_pod=False,
termination_grace_period=0,
)
def test_envs_from_configmaps(self, mock_client, mock_monitor, mock_start):
# GIVEN
configmap = 'test-configmap'
# WHEN
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
configmaps=[configmap],
)
# THEN
mock_monitor.return_value = (State.SUCCESS, None)
context = create_context(k)
k.execute(context)
assert mock_start.call_args[0][0].spec.containers[0].env_from == [
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap))
]
def test_pod_resources(self):
resources = {
'limit_cpu': 0.25,
'limit_memory': '64Mi',
'limit_ephemeral_storage': '2Gi',
'request_cpu': '250m',
'request_memory': '64Mi',
'request_ephemeral_storage': '1Gi',
}
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
resources=resources,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['resources'] = {
'requests': {
'memory': '64Mi',
'cpu': '250m',
'ephemeral-storage': '1Gi'
},
'limits': {
'memory': '64Mi',
'cpu': 0.25,
'ephemeral-storage': '2Gi'
},
}
assert self.expected_pod == actual_pod
def test_envs_from_configmaps(self, mock_client, mock_monitor, mock_start):
# GIVEN
from airflow.utils.state import State
configmap_name = "test-config-map"
env_from = [k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap_name))]
# WHEN
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,
env_from=env_from,
)
# THEN
mock_monitor.return_value = (State.SUCCESS, None)
context = create_context(k)
k.execute(context)
self.assertEqual(mock_start.call_args[0][0].spec.containers[0].env_from, env_from)
def test_env_vars(self):
# WHEN
env_vars = [k8s.V1EnvVar(name="{{ bar }}", value='{{ foo }}')]
from tests.models import DEFAULT_DATE
with DAG("test-dag", start_date=DEFAULT_DATE):
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
env_vars=env_vars,
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
k.render_template_fields(context={
"foo": "footemplated",
"bar": "bartemplated"
})
assert k.env_vars[0].value == "footemplated"
assert k.env_vars[0].name == "bartemplated"
def test_config_path(self, client_mock, launcher_mock):
from airflow.utils.state import State
file_path = "/tmp/fake_file"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
config_file=file_path,
cluster_context='default',
)
launcher_mock.return_value = (State.SUCCESS, None)
k.execute(None)
client_mock.assert_called_once_with(
in_cluster=False,
cluster_context='default',
config_file=file_path,
)
def test_xcom_push(self):
return_value = '{"foo": "bar"\n, "buzz": 2}'
args = ['echo \'{}\' > /airflow/xcom/return.json'.format(return_value)]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=True,
)
self.assertEqual(k.execute(None), json.loads(return_value))
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
volume = self.api_client.sanitize_for_serialization(PodDefaults.VOLUME)
volume_mount = self.api_client.sanitize_for_serialization(PodDefaults.VOLUME_MOUNT)
container = self.api_client.sanitize_for_serialization(PodDefaults.SIDECAR_CONTAINER)
self.expected_pod['spec']['containers'][0]['args'] = args
self.expected_pod['spec']['containers'][0]['volumeMounts'].insert(0, volume_mount)
self.expected_pod['spec']['volumes'].insert(0, volume)
self.expected_pod['spec']['containers'].append(container)
self.assertEqual(self.expected_pod, actual_pod)
def test_config_path(self):
file_path = "/tmp/fake_file"
k = KubernetesPodOperator(
namespace="default",
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
config_file=file_path,
cluster_context="default",
)
self.monitor_mock.return_value = (State.SUCCESS, None)
self.client_mock.list_namespaced_pod.return_value = []
context = self.create_context(k)
k.execute(context=context)
self.client_mock.assert_called_once_with(
in_cluster=False,
cluster_context="default",
config_file=file_path,
)
def test_pod_delete_even_on_launcher_error(
self,
mock_client,
delete_pod_mock,
monitor_pod_mock,
start_pod_mock):
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
is_delete_operator_pod=True,
)
monitor_pod_mock.side_effect = AirflowException('fake failure')
with self.assertRaises(AirflowException):
context = self.create_context(k)
k.execute(context=context)
assert delete_pod_mock.called
def test_envs_from_secrets(self, mock_client, monitor_mock, start_mock):
# GIVEN
secret_ref = 'secret_name'
secrets = [Secret('env', None, secret_ref)]
# WHEN
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
secrets=secrets,
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
# THEN
monitor_mock.return_value = (State.SUCCESS, None)
context = create_context(k)
k.execute(context)
assert start_mock.call_args[0][0].spec.containers[0].env_from == [
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name=secret_ref))
]
def test_config_path_move(self):
new_config_path = '/tmp/kube_config'
old_config_path = get_kubeconfig_path()
shutil.copy(old_config_path, new_config_path)
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test1",
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
is_delete_operator_pod=False,
config_file=new_config_path,
)
context = create_context(k)
k.execute(context)
expected_pod = copy(self.expected_pod)
expected_pod['metadata']['labels']['already_checked'] = 'True'
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
assert expected_pod == actual_pod
def test_push_xcom_pod_info(self):
k = KubernetesPodOperator(
namespace="default",
image="ubuntu:16.04",
cmds=["bash", "-cx"],
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
pod = self.run_pod(k)
ti = TaskInstance(task=k, execution_date=DEFAULT_DATE)
pod_name = ti.xcom_pull(task_ids=k.task_id, key='pod_name')
pod_namespace = ti.xcom_pull(task_ids=k.task_id, key='pod_namespace')
assert pod_name and pod_name == pod.metadata.name
assert pod_namespace and pod_namespace == pod.metadata.namespace
def test_reattach_failing_pod_once(self):
from airflow.utils.state import State
client = kube_client.get_kube_client(in_cluster=False)
name = "test"
namespace = "default"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["exit 1"],
labels={"foo": "bar"},
name="test",
task_id=name,
in_cluster=False,
do_xcom_push=False,
is_delete_operator_pod=False,
termination_grace_period=0,
)
context = create_context(k)
with mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.monitor_pod"
) as monitor_mock:
monitor_mock.return_value = (State.SUCCESS, None, None)
k.execute(context)
name = k.pod.metadata.name
pod = client.read_namespaced_pod(name=name, namespace=namespace)
while pod.status.phase != "Failed":
pod = client.read_namespaced_pod(name=name, namespace=namespace)
with pytest.raises(AirflowException):
k.execute(context)
pod = client.read_namespaced_pod(name=name, namespace=namespace)
assert pod.metadata.labels["already_checked"] == "True"
with mock.patch(
"airflow.providers.cncf.kubernetes"
".operators.kubernetes_pod.KubernetesPodOperator"
".create_new_pod_for_operator"
) as create_mock:
create_mock.return_value = ("success", {}, {})
k.execute(context)
create_mock.assert_called_once()
def test_do_xcom_push_defaults_false(self):
new_config_path = '/tmp/kube_config'
old_config_path = get_kubeconfig_path()
shutil.copy(old_config_path, new_config_path)
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,
config_file=new_config_path,
)
self.assertFalse(k.do_xcom_push)
def test_do_xcom_push_defaults_false(self):
new_config_path = '/tmp/kube_config'
old_config_path = os.path.expanduser('~/.kube/config')
shutil.copy(old_config_path, new_config_path)
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
config_file=new_config_path,
)
self.assertFalse(k.do_xcom_push)
def test_full_pod_spec(self):
pod_spec = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name="hello",
labels={"foo": "bar"},
namespace="mynamespace"),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(
name="base",
image="ubuntu:16.04",
command=["something"],
)
]),
)
k = KubernetesPodOperator(
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context="default",
full_pod_spec=pod_spec,
)
pod = k.create_pod_request_obj()
assert pod.metadata.name == pod_spec.metadata.name
assert pod.metadata.labels == pod_spec.metadata.labels
assert pod.metadata.namespace == pod_spec.metadata.namespace
assert pod.spec.containers[0].image == pod_spec.spec.containers[
0].image
assert pod.spec.containers[0].command == pod_spec.spec.containers[
0].command
# kwargs take precedence, however
image = "some.custom.image:andtag"
name_base = "world"
k = KubernetesPodOperator(
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context="default",
full_pod_spec=pod_spec,
name=name_base,
image=image,
)
pod = k.create_pod_request_obj()
# make sure the kwargs takes precedence (and that name is randomized)
assert pod.metadata.name.startswith(name_base)
assert pod.metadata.name != name_base
assert pod.spec.containers[0].image == image
