def set_labels_node(self, nodename, labels, kube_config=None):
"""Set label for a node
Parameters
----------
nodename: str
hostname of the node
labels: string
the new k8s labels used to label this node, the format is: key1=value1,key2=value2,...
kube_config: kubernetes.client.configuration.Configuration
the configuration to the kubernetes cluster
Returns
-------
V1Node (https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1Node.md)
if label node successfully
None
if label node unsuccessfully
"""
if kube_config:
api_client = ApiClient(kube_config)
else:
api_client = ApiClient()
v1 = client.CoreV1Api(api_client)
logger.debug('Label node %s with %s' % (nodename, labels))
for node in v1.list_node().items:
if node.metadata.name == nodename:
break
else:
logger.warning('Node %s does not exist' % nodename)
return None
body = {
'metadata': {
'labels': dict()
}
}
for label in labels.strip().split(','):
key, val = label.strip().split('=')
body['metadata']['labels'][key] = val
return v1.patch_node(name=nodename, body=body)
def setUp(self):
self.maxDiff = None
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': mock.ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-'),
'execution_date': '2016-01-01T0100000100-a2f50a31f',
'dag_id': 'dag',
'task_id': 'task',
'try_number': '1',
},
},
'spec': {
'affinity': {},
'containers': [
{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'envFrom': [],
'resources': {},
'name': 'base',
'ports': [],
'volumeMounts': [],
}
],
'hostNetwork': False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy': 'Never',
'securityContext': {},
'tolerations': [],
'volumes': [],
},
}
def deserialize_model_file(path: str) -> k8s.V1Pod:
"""
:param path: Path to the file
:return: a kubernetes.client.models.V1Pod
Unfortunately we need access to the private method
``_ApiClient__deserialize_model`` from the kubernetes client.
This issue is tracked here; https://github.com/kubernetes-client/python/issues/977.
"""
api_client = ApiClient()
if os.path.exists(path):
with open(path) as stream:
pod = yaml.safe_load(stream)
else:
pod = yaml.safe_load(path)
# pylint: disable=protected-access
return api_client._ApiClient__deserialize_model(pod, k8s.V1Pod)
def setUp(self):
if AirflowKubernetesScheduler is None:
self.skipTest("kubernetes python package is not installed")
self.kube_config = mock.MagicMock()
self.kube_config.airflow_home = '/'
self.kube_config.airflow_dags = 'dags'
self.kube_config.airflow_logs = 'logs'
self.kube_config.dags_volume_subpath = None
self.kube_config.logs_volume_subpath = None
self.kube_config.dags_in_image = False
self.kube_config.dags_folder = None
self.kube_config.git_dags_folder_mount_point = None
self.kube_config.kube_labels = {
'dag_id': 'original_dag_id',
'my_label': 'label_id'
}
self.api_client = ApiClient()
def generate_pod_yaml(args):
"""Generates yaml files for each task in the DAG. Used for testing output of KubernetesExecutor"""
from kubernetes.client.api_client import ApiClient
from airflow.executors.kubernetes_executor import AirflowKubernetesScheduler, KubeConfig
from airflow.kubernetes import pod_generator
from airflow.kubernetes.pod_generator import PodGenerator
from airflow.settings import pod_mutation_hook
execution_date = args.execution_date
dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
yaml_output_path = args.output_path
kube_config = KubeConfig()
for task in dag.tasks:
ti = TaskInstance(task, execution_date)
pod = PodGenerator.construct_pod(
dag_id=args.dag_id,
task_id=ti.task_id,
pod_id=AirflowKubernetesScheduler._create_pod_id( # pylint: disable=W0212
args.dag_id, ti.task_id),
try_number=ti.try_number,
kube_image=kube_config.kube_image,
date=ti.execution_date,
command=ti.command_as_list(),
pod_override_object=PodGenerator.from_obj(ti.executor_config),
worker_uuid="worker-config",
namespace=kube_config.executor_namespace,
base_worker_pod=PodGenerator.deserialize_model_file(
kube_config.pod_template_file))
pod_mutation_hook(pod)
api_client = ApiClient()
date_string = pod_generator.datetime_to_label_safe_datestring(
execution_date)
yaml_file_name = f"{args.dag_id}_{ti.task_id}_{date_string}.yml"
os.makedirs(os.path.dirname(yaml_output_path +
"/airflow_yaml_output/"),
exist_ok=True)
with open(yaml_output_path + "/airflow_yaml_output/" + yaml_file_name,
"w") as output:
sanitized_pod = api_client.sanitize_for_serialization(pod)
output.write(yaml.dump(sanitized_pod))
print(
f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
def deploy_k8s_resources(self,
path=None,
files=None,
kube_config=None,
namespace="default"):
"""Deploy k8s resources on a k8s cluster from deployment yaml files
Parameters
----------
path: str
the path to the directory which stores all the deployment yaml files
files: list
a list of yaml files to use for deployment
kube_config: kubernetes.client.configuration.Configuration
the configuration to the kubernetes cluster
namespace: str
a namespace for k8s working with
"""
if kube_config:
api_client = ApiClient(kube_config)
else:
api_client = ApiClient()
if path is not None:
files = list()
for file in os.listdir(path):
if file.endswith('.yaml'):
files.append(os.path.join(path, file))
for file in files:
logger.info('--> Deploying file %s' % file.split('/')[-1])
try:
utils.create_from_yaml(k8s_client=api_client,
yaml_file=file,
namespace=namespace)
logger.debug('Deploy file %s successfully' % file)
except FailToCreateError as e:
for api_exception in e.api_exceptions:
body = json.loads(api_exception.body)
logger.error('Error: %s, because: %s' %
(api_exception.reason, body['message']))
def setUp(self):
self.maxDiff = None # pylint: disable=invalid-name
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-')
}
},
'spec': {
'affinity': {},
'containers': [{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'imagePullPolicy': 'IfNotPresent',
'envFrom': [],
'name': 'base',
'ports': [],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy':
'Never',
'securityContext': {},
'serviceAccountName':
'default',
'tolerations': [],
'volumes': [],
}
}
def __init__(self, *args, **kwargs):
self.networking_api = kubernetes.client.NetworkingV1beta1Api()
self.api: CoreV1Api = kubernetes.client.CoreV1Api()
self.secret_creator = TokenSecretCreator()
self.api_client: ApiClient = ApiClient()
self.custom_objects_api: CustomObjectsApi = CustomObjectsApi(
self.api_client)
super(KubernetesGameManager, self).__init__(*args, **kwargs)
self._create_ingress_paths_for_existing_games()
def _get_k8s_api_client() -> ApiClient:
# Create new client everytime to avoid token refresh issues
# https://github.com/kubernetes-client/python/issues/741
# https://github.com/kubernetes-client/python-base/issues/125
if bool(
util.strtobool(
os.environ.get('LOAD_IN_CLUSTER_KUBECONFIG', 'false'))):
config.load_incluster_config()
return ApiClient()
return config.new_client_from_config()
def execute_command(self, pod_name, command, kube_config=None, kube_namespace='default'):
"""Execute a command on a pod
Parameters
----------
pod_name: string
the name of the pod to run the command
command: string
the command to run on the pod
kube_config: kubernetes.client.configuration.Configuration
the configuration to the kubernetes cluster
kube_namespace: string
the k8s namespace to perform the k8s resources operation on,
the default namespace is 'default'
Returns
-------
string
the std output when run the command in the pod
"""
if kube_config:
api_client = ApiClient(kube_config)
else:
api_client = ApiClient()
v1 = client.CoreV1Api(api_client)
logger.debug('Run command %s on pod %s' % (command, pod_name))
if ' ' in command:
command = _split_argument(command)
return stream(v1.connect_get_namespaced_pod_exec,
pod_name,
namespace=kube_namespace,
command=command,
stderr=True,
stdin=False,
stdout=True,
tty=False,
_preload_content=True)
def api_init(kube_config_file=None, host=None, token_filename=None, cert_filename=None, context=None):
global CoreV1Api
global RbacAuthorizationV1Api
global api_temp
if host and token_filename:
# remotely
token_filename = os.path.abspath(token_filename)
if cert_filename:
cert_filename = os.path.abspath(cert_filename)
BearerTokenLoader(host=host, token_filename=token_filename, cert_filename=cert_filename).load_and_set()
CoreV1Api = client.CoreV1Api()
RbacAuthorizationV1Api = client.RbacAuthorizationV1Api()
api_temp = ApiClientTemp()
elif kube_config_file:
config.load_kube_config(os.path.abspath(kube_config_file))
CoreV1Api = client.CoreV1Api()
RbacAuthorizationV1Api = client.RbacAuthorizationV1Api()
api_from_config = config.new_client_from_config(kube_config_file)
api_temp = ApiClientTemp(configuration=api_from_config.configuration)
else:
configuration = Configuration()
api_client = ApiClient()
if running_in_docker_container():
# TODO: Consider using config.load_incluster_config() from container created by Kubernetes. Required service account with privileged permissions.
# Must have mounted volume
container_volume_prefix = '/tmp'
kube_config_bak_path = '/KubiScan/config_bak'
if not os.path.isfile(kube_config_bak_path):
copyfile(container_volume_prefix + os.path.expandvars('$CONF_PATH'), kube_config_bak_path)
replace(kube_config_bak_path, ': /', ': /tmp/')
config.load_kube_config(kube_config_bak_path, context=context, client_configuration=configuration)
else:
config.load_kube_config(context=context, client_configuration=configuration)
api_client = ApiClient(configuration=configuration)
CoreV1Api = client.CoreV1Api(api_client=api_client)
RbacAuthorizationV1Api = client.RbacAuthorizationV1Api(api_client=api_client)
api_temp = ApiClientTemp(configuration=configuration)
def _get_client_with_patched_configuration(cfg: Optional[Configuration]) -> client.CoreV1Api:
"""
This is a workaround for supporting api token refresh in k8s client.
The function can be replace with `return client.CoreV1Api()` once the
upstream client supports token refresh.
"""
if cfg:
return client.CoreV1Api(api_client=ApiClient(configuration=cfg))
else:
return client.CoreV1Api()
def client(self) -> ApiClient:
"""Get Kubernetes client configured from kubeconfig"""
config = Configuration()
try:
if self.local:
load_incluster_config(client_configuration=config)
else:
load_kube_config_from_dict(self.kubeconfig,
client_configuration=config)
return ApiClient(config)
except ConfigException as exc:
raise ServiceConnectionInvalid from exc
def _convert_to_airflow_pod(pod):
"""
Converts a k8s V1Pod object into an `airflow.kubernetes.pod.Pod` object.
This function is purely for backwards compatibility
"""
base_container = pod.spec.containers[0] # type: k8s.V1Container
env_vars, secrets = _extract_env_vars_and_secrets(base_container.env)
volumes = _extract_volumes(pod.spec.volumes)
api_client = ApiClient()
init_containers = pod.spec.init_containers
image_pull_secrets = pod.spec.image_pull_secrets or []
if pod.spec.init_containers is not None:
init_containers = [
api_client.sanitize_for_serialization(i)
for i in pod.spec.init_containers
]
dummy_pod = Pod(
image=base_container.image,
envs=env_vars,
cmds=base_container.command,
args=base_container.args,
labels=pod.metadata.labels,
annotations=pod.metadata.annotations,
node_selectors=pod.spec.node_selector,
name=pod.metadata.name,
ports=_extract_ports(base_container.ports),
volumes=volumes,
volume_mounts=_extract_volume_mounts(base_container.volume_mounts),
namespace=pod.metadata.namespace,
image_pull_policy=base_container.image_pull_policy or 'IfNotPresent',
tolerations=pod.spec.tolerations,
init_containers=init_containers,
image_pull_secrets=",".join([i.name for i in image_pull_secrets]),
resources=base_container.resources,
service_account_name=pod.spec.service_account_name,
secrets=secrets,
affinity=api_client.sanitize_for_serialization(pod.spec.affinity),
hostnetwork=pod.spec.host_network,
security_context=_extract_security_context(pod.spec.security_context))
return dummy_pod
def delete_namespace(self, namespace=None, kube_config=None):
"""Delete a namespace from a k8s cluster
Parameters
----------
namespace: str
a namespace for k8s working with
kube_config: kubernetes.client.configuration.Configuration
the configuration to the kubernetes cluster
Returns
-------
bool
True: delete successfully
False: delete unsuccessfully
"""
if kube_config:
api_client = ApiClient(kube_config)
else:
api_client = ApiClient()
v1 = client.CoreV1Api(api_client)
logger.debug('Deleting namespace %s' % namespace)
for ns in v1.list_namespace().items:
if ns.metadata.name == namespace:
v1.delete_namespace(name=namespace)
logger.debug('Waiting for namespace %s to be deleted' %
namespace)
for i in range(100):
for ns in v1.list_namespace().items:
if ns.metadata.name == namespace:
sleep(5)
break
else:
return True
else:
logger.warning('Namespace %s does not exist' % namespace)
return False
def main():
config.load_kube_config()
client = DynamicClient(ApiClient())
ret = {}
for resource in client._resources:
if resource.namespaced:
key = resource.urls['namespaced_full']
else:
key = resource.urls['full']
ret[key] = {k: v for k, v in resource.__dict__.items() if k != 'client'}
print(yaml.safe_dump(ret))
return 0
def job_to_dict(job: V1Job) -> Optional[dict]:
"""convert V1Job to dictionary
Note that this is *different* than what is returned by V1Job.to_dict().
That method uses object attribute names, which are often different than
the names used in the REST api and yaml/json spec files.
Args:
job: V1Job instance
Returns:
dictionary representation on success, None otherwise
"""
return ApiClient().sanitize_for_serialization(job)
def main():
config.load_kube_config()
client = DynamicClient(ApiClient())
ret = {}
for resource in client.resources:
if resource.namespaced:
key = resource.urls['namespaced_full']
else:
key = resource.urls['full']
ret[key] = {k: v for k, v in resource.__dict__.items() if k not in ('client', 'subresources')}
ret[key]['subresources'] = {}
for name, value in resource.subresources.items():
ret[key]['subresources'][name] = {k: v for k, v in value.__dict__.items() if k != 'parent'}
print(yaml.safe_dump(ret))
return 0
def get_cli(cluster_name):
global _client_cache
if cluster_name in _client_cache:
return _client_cache.get(cluster_name, None)
cluster = Cluster.query.filter_by(name=cluster_name).first()
if cluster is None:
return None
cfg = Configuration()
# cfg.debug = True
cfg.host = cluster.addr
cfg.ssl_ca_cert = _create_temp_file_with_content(
base64.b64decode(cluster.cert))
cfg.api_key['authorization'] = "Bearer {}".format(cluster.access_token)
api = ApiClient(cfg)
_client_cache[cluster_name] = api
return api
class HelmTemplate:
output: Optional[str] = None
model: Optional[Any] = None
name: str = "RELEASE-NAME"
api_client: ApiClient = ApiClient()
def render(self, values: DagsterHelmValues) -> List[Any]:
with NamedTemporaryFile() as tmp_file:
values_json = json.loads(
values.json(exclude_none=True, by_alias=True))
pprint(values_json)
content = yaml.dump(values_json)
tmp_file.write(content.encode())
tmp_file.flush()
command = [
"helm",
"template",
self.name,
os.path.join(git_repo_root(), "helm", "dagster"),
"--debug",
*['--values', tmp_file.name],
]
if self.output:
command += ["--show-only", self.output]
templates = subprocess.check_output(command)
print("\n--- Helm Templates ---") # pylint: disable=print-call
print(templates.decode()) # pylint: disable=print-call
k8s_objects = [
k8s_object for k8s_object in yaml.full_load_all(templates)
if k8s_object
]
if self.model:
k8s_objects = [
self.api_client._ApiClient__deserialize_model( # pylint: disable=W0212
k8s_object, self.model) for k8s_object in k8s_objects
]
return k8s_objects
def delete_vnf(self, uuid, tenant=None):
if uuid:
todelete_uuids = [uuid]
else:
todelete_uuids = [
uuid for uuid in self.vnfs if self.vnfs[uuid].tenant == tenant
]
for todelete_uuid in todelete_uuids:
vnf = self.vnfs[todelete_uuid]
handle_yaml(k8s_client=ApiClient(),
yaml_file=vnf.get_k8s_desc_file().name,
mode="delete",
namespace=tenant)
del self.vnfs[todelete_uuid]
if not [vnf for vnf in self.vnfs.values() if vnf.tenant == tenant]:
v1 = CoreV1Api()
try:
v1.delete_namespace(tenant)
except:
pass
def create_vnf(self, uuid, **params):
v1 = CoreV1Api()
namespaces = [ns.metadata.name for ns in v1.list_namespace().items]
tenant = params["tenant"]
if tenant not in namespaces:
ns_file = open("/etc/lightmano/k8s/_internal/namespace.yaml")
ns = ns_file.read()
ns_file.close()
ns = ns.replace("-NAME-", tenant)
ns = yaml.safe_load(ns)
v1.create_namespace(ns)
rbac_api = RbacAuthorizationV1Api()
role_file = open("/etc/lightmano/k8s/_internal/ns_role.yaml")
role = role_file.read()
role_file.close()
role = yaml.safe_load(role)
rbac_api.create_namespaced_role(namespace=tenant, body=role)
rolebind_file = open(
"/etc/lightmano/k8s/_internal/ns_role_binding.yaml")
rolebinding = rolebind_file.read()
rolebind_file.close()
rolebinding = yaml.safe_load(rolebinding)
rbac_api.create_namespaced_role_binding(namespace=tenant,
body=rolebinding)
vnf = VNF(uuid, **params)
handle_yaml(k8s_client=ApiClient(),
yaml_file=vnf.get_k8s_desc_file().name,
mode="create",
namespace=tenant)
self.vnfs[uuid] = vnf
def __init__(self, api_client=None):
if api_client is None:
api_client = ApiClient()
self.api_client = api_client
class TestKubernetesPodOperatorSystem(unittest.TestCase):
def get_current_task_name(self):
# reverse test name to make pod name unique (it has limited length)
return "_" + unittest.TestCase.id(self).replace(".", "_")[::-1]
def setUp(self):
self.maxDiff = None
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': mock.ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-'),
'execution_date': '2016-01-01T0100000100-a2f50a31f',
'dag_id': 'dag',
'task_id': 'task',
'try_number': '1',
},
},
'spec': {
'affinity': {},
'containers': [{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'envFrom': [],
'resources': {},
'name': 'base',
'ports': [],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy':
'Never',
'securityContext': {},
'tolerations': [],
'volumes': [],
},
}
def tearDown(self):
client = kube_client.get_kube_client(in_cluster=False)
client.delete_collection_namespaced_pod(namespace="default")
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.start_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.monitor_pod"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_image_pull_secrets_correctly_set(self, mock_client, monitor_mock,
start_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',
)
monitor_mock.return_value = (State.SUCCESS, None, None)
context = create_context(k)
k.execute(context=context)
assert start_mock.call_args[0][0].spec.image_pull_secrets == [
k8s.V1LocalObjectReference(name=fake_pull_secrets)
]
def test_working_pod(self):
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,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
assert self.expected_pod['spec'] == actual_pod['spec']
assert self.expected_pod['metadata']['labels'] == actual_pod[
'metadata']['labels']
def test_pod_node_selectors(self):
node_selectors = {'beta.kubernetes.io/os': 'linux'}
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,
node_selectors=node_selectors,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['nodeSelector'] = node_selectors
assert self.expected_pod == actual_pod
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_pod_affinity(self):
affinity = {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': {
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'beta.kubernetes.io/os',
'operator': 'In',
'values': ['linux']
}]
}]
}
}
}
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,
affinity=affinity,
)
context = create_context(k)
k.execute(context=context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['affinity'] = affinity
assert self.expected_pod == actual_pod
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 = 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_volume_mount(self):
with patch.object(PodLauncher, 'log') as mock_logger:
volume_mount = VolumeMount('test-volume',
mount_path='/tmp/test_volume',
sub_path=None,
read_only=False)
volume_config = {
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}
volume = Volume(name='test-volume', configs=volume_config)
args = [
"echo \"retrieved from mount\" > /tmp/test_volume/test.txt "
"&& cat /tmp/test_volume/test.txt"
]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
volume_mounts=[volume_mount],
volumes=[volume],
is_delete_operator_pod=False,
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context=context)
mock_logger.info.assert_any_call('retrieved from mount')
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['args'] = args
self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{
'name':
'test-volume',
'mountPath':
'/tmp/test_volume',
'readOnly':
False
}]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
assert 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",
task_id="task",
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_run_as_user_non_root(self):
security_context = {
'securityContext': {
'runAsUser': 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,
)
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_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,
)
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_faulty_service_account(self):
bad_service_account_name = "foobar"
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,
startup_timeout_seconds=5,
service_account_name=bad_service_account_name,
)
with pytest.raises(ApiException):
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec'][
'serviceAccountName'] = bad_service_account_name
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",
task_id="task",
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_xcom_push(self):
return_value = '{"foo": "bar"\n, "buzz": 2}'
args = [f'echo \'{return_value}\' > /airflow/xcom/return.json']
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,
)
context = create_context(k)
assert k.execute(context) == 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)
assert self.expected_pod == actual_pod
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.start_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.monitor_pod"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
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, 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))
]
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.start_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_launcher.PodLauncher.monitor_pod"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
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, 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_env_vars(self):
# WHEN
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
env_vars={
"ENV1": "val1",
"ENV2": "val2",
},
pod_runtime_info_envs=[PodRuntimeInfoEnv("ENV3", "status.podIP")],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context)
# THEN
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['env'] = [
{
'name': 'ENV1',
'value': 'val1'
},
{
'name': 'ENV2',
'value': 'val2'
},
{
'name': 'ENV3',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
},
]
assert self.expected_pod == actual_pod
def test_pod_template_file_with_overrides_system(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
labels={
"foo": "bar",
"fizz": "buzz"
},
env_vars={"env_name": "value"},
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
assert result is not None
assert k.pod.metadata.labels == {
'fizz': 'buzz',
'foo': 'bar',
'airflow_version': mock.ANY,
'dag_id': 'dag',
'execution_date': mock.ANY,
'kubernetes_pod_operator': 'True',
'task_id': mock.ANY,
'try_number': '1',
}
assert k.pod.spec.containers[0].env == [
k8s.V1EnvVar(name="env_name", value="value")
]
assert result == {"hello": "world"}
def test_init_container(self):
# GIVEN
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=volume_mounts,
command=["bash", "-cx"],
args=["echo 10"],
)
volume_config = {'persistentVolumeClaim': {'claimName': 'test-volume'}}
volume = Volume(name='test-volume', configs=volume_config)
expected_init_container = {
'name':
'init-container',
'image':
'ubuntu:16.04',
'command': ['bash', '-cx'],
'args': ['echo 10'],
'env': [{
'name': 'key1',
'value': 'value1'
}, {
'name': 'key2',
'value': 'value2'
}],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'test-volume',
'readOnly': True
}],
}
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
volumes=[volume],
init_containers=[init_container],
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['initContainers'] = [expected_init_container]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
assert self.expected_pod == actual_pod
class TestKubernetesPodOperatorSystem(unittest.TestCase):
def get_current_task_name(self):
# reverse test name to make pod name unique (it has limited length)
return "_" + unittest.TestCase.id(self).replace(".", "_")[::-1]
def setUp(self):
self.maxDiff = None # pylint: disable=invalid-name
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-'),
'execution_date': '2016-01-01T0100000100-a2f50a31f',
'dag_id': 'dag',
'task_id': ANY,
'try_number': '1',
},
},
'spec': {
'affinity': {},
'containers': [{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'envFrom': [],
'resources': {},
'name': 'base',
'ports': [],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'restartPolicy':
'Never',
'securityContext': {},
'serviceAccountName':
'default',
'tolerations': [],
'volumes': [],
},
}
def tearDown(self) -> None:
client = kube_client.get_kube_client(in_cluster=False)
client.delete_collection_namespaced_pod(namespace="default")
import time
time.sleep(1)
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_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,
config_file=new_config_path,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod, actual_pod)
def test_working_pod(self):
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,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod['spec'], actual_pod['spec'])
self.assertEqual(self.expected_pod['metadata']['labels'],
actual_pod['metadata']['labels'])
def test_delete_operator_pod(self):
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,
is_delete_operator_pod=True,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod['spec'], actual_pod['spec'])
self.assertEqual(self.expected_pod['metadata']['labels'],
actual_pod['metadata']['labels'])
def test_pod_hostnetwork(self):
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,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['hostNetwork'] = True
self.assertEqual(self.expected_pod['spec'], actual_pod['spec'])
self.assertEqual(self.expected_pod['metadata']['labels'],
actual_pod['metadata']['labels'])
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
self.assertEqual(self.expected_pod['spec'], actual_pod['spec'])
self.assertEqual(self.expected_pod['metadata']['labels'],
actual_pod['metadata']['labels'])
def test_pod_schedulername(self):
scheduler_name = "default-scheduler"
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,
schedulername=scheduler_name,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['schedulerName'] = scheduler_name
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_node_selectors(self):
node_selectors = {'beta.kubernetes.io/os': 'linux'}
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,
node_selectors=node_selectors,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['nodeSelector'] = node_selectors
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_resources(self):
resources = k8s.V1ResourceRequirements(
requests={
'memory': '64Mi',
'cpu': '250m',
'ephemeral-storage': '1Gi'
},
limits={
'memory': '64Mi',
'cpu': 0.25,
'nvidia.com/gpu': None,
'ephemeral-storage': '2Gi'
},
)
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,
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,
'nvidia.com/gpu': None,
'ephemeral-storage': '2Gi'
},
}
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_affinity(self):
affinity = {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': {
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'beta.kubernetes.io/os',
'operator': 'In',
'values': ['linux']
}]
}]
}
}
}
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,
affinity=affinity,
)
context = create_context(k)
k.execute(context=context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['affinity'] = affinity
self.assertEqual(self.expected_pod, actual_pod)
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_volume_mount(self):
with mock.patch.object(PodLauncher, 'log') as mock_logger:
volume_mount = k8s.V1VolumeMount(name='test-volume',
mount_path='/tmp/test_volume',
sub_path=None,
read_only=False)
volume = k8s.V1Volume(
name='test-volume',
persistent_volume_claim=k8s.
V1PersistentVolumeClaimVolumeSource(claim_name='test-volume'),
)
args = [
"echo \"retrieved from mount\" > /tmp/test_volume/test.txt "
"&& cat /tmp/test_volume/test.txt"
]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
volume_mounts=[volume_mount],
volumes=[volume],
name="test-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context=context)
mock_logger.info.assert_any_call('retrieved from mount')
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['args'] = args
self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{
'name':
'test-volume',
'mountPath':
'/tmp/test_volume',
'readOnly':
False
}]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
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
self.assertEqual(self.expected_pod, actual_pod)
def test_run_as_user_non_root(self):
security_context = {
'securityContext': {
'runAsUser': 1000,
}
}
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
self.assertEqual(self.expected_pod, actual_pod)
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-fs-group",
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
self.assertEqual(self.expected_pod, actual_pod)
def test_faulty_image(self):
bad_image_name = "foobar"
k = KubernetesPodOperator(
namespace='default',
image=bad_image_name,
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,
startup_timeout_seconds=5,
)
with self.assertRaises(AirflowException):
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0][
'image'] = bad_image_name
self.assertEqual(self.expected_pod, actual_pod)
def test_faulty_service_account(self):
bad_service_account_name = "foobar"
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,
startup_timeout_seconds=5,
service_account_name=bad_service_account_name,
)
with self.assertRaises(ApiException):
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec'][
'serviceAccountName'] = bad_service_account_name
self.assertEqual(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 self.assertRaises(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
self.assertEqual(self.expected_pod, actual_pod)
def test_xcom_push(self):
return_value = '{"foo": "bar"\n, "buzz": 2}'
args = [f'echo \'{return_value}\' > /airflow/xcom/return.json']
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
name="test-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=True,
)
context = create_context(k)
self.assertEqual(k.execute(context), 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) # noqa
self.expected_pod['spec']['volumes'].insert(0, volume)
self.expected_pod['spec']['containers'].append(container)
self.assertEqual(self.expected_pod, actual_pod)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.start_pod")
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
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)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.start_pod")
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_envs_from_secrets(self, mock_client, monitor_mock, start_mock):
# GIVEN
from airflow.utils.state import State
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-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
)
# THEN
monitor_mock.return_value = (State.SUCCESS, None)
context = create_context(k)
k.execute(context)
self.assertEqual(
start_mock.call_args[0][0].spec.containers[0].env_from,
[
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(
name=secret_ref))
],
)
def test_env_vars(self):
# WHEN
env_vars = [
k8s.V1EnvVar(name="ENV1", value="val1"),
k8s.V1EnvVar(name="ENV2", value="val2"),
k8s.V1EnvVar(
name="ENV3",
value_from=k8s.V1EnvVarSource(
field_ref=k8s.V1ObjectFieldSelector(
field_path="status.podIP")),
),
]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
env_vars=env_vars,
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,
)
context = create_context(k)
k.execute(context)
# THEN
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['env'] = [
{
'name': 'ENV1',
'value': 'val1'
},
{
'name': 'ENV2',
'value': 'val2'
},
{
'name': 'ENV3',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
},
]
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_template_file_system(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
self.assertIsNotNone(result)
self.assertDictEqual(result, {"hello": "world"})
def test_pod_template_file_with_overrides_system(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
labels={
"foo": "bar",
"fizz": "buzz"
},
env_vars=[k8s.V1EnvVar(name="env_name", value="value")],
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
self.assertIsNotNone(result)
self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'})
self.assertEqual(k.pod.spec.containers[0].env,
[k8s.V1EnvVar(name="env_name", value="value")])
self.assertDictEqual(result, {"hello": "world"})
def test_pod_template_file_with_full_pod_spec(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
pod_spec = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(labels={
"foo": "bar",
"fizz": "buzz"
}, ),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(
name="base",
env=[k8s.V1EnvVar(name="env_name", value="value")],
)
]),
)
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
pod_template_file=fixture,
full_pod_spec=pod_spec,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
self.assertIsNotNone(result)
self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'})
self.assertEqual(k.pod.spec.containers[0].env,
[k8s.V1EnvVar(name="env_name", value="value")])
self.assertDictEqual(result, {"hello": "world"})
def test_full_pod_spec(self):
pod_spec = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(labels={
"foo": "bar",
"fizz": "buzz"
},
namespace="default",
name="test-pod"),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
image="perl",
command=["/bin/bash"],
args=[
"-c",
'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json'
],
env=[k8s.V1EnvVar(name="env_name", value="value")],
)
],
restart_policy="Never",
),
)
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
full_pod_spec=pod_spec,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
self.assertIsNotNone(result)
self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'})
self.assertEqual(k.pod.spec.containers[0].env,
[k8s.V1EnvVar(name="env_name", value="value")])
self.assertDictEqual(result, {"hello": "world"})
def test_init_container(self):
# GIVEN
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=volume_mounts,
command=["bash", "-cx"],
args=["echo 10"],
)
volume = k8s.V1Volume(
name='test-volume',
persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(
claim_name='test-volume'),
)
expected_init_container = {
'name':
'init-container',
'image':
'ubuntu:16.04',
'command': ['bash', '-cx'],
'args': ['echo 10'],
'env': [{
'name': 'key1',
'value': 'value1'
}, {
'name': 'key2',
'value': 'value2'
}],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'test-volume',
'readOnly': True
}],
}
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(),
volumes=[volume],
init_containers=[init_container],
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['initContainers'] = [expected_init_container]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
self.assertEqual(self.expected_pod, actual_pod)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.start_pod")
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_template_file(
self,
mock_client,
monitor_mock,
start_mock # pylint: disable=unused-argument
):
from airflow.utils.state import State
path = sys.path[0] + '/tests/kubernetes/pod.yaml'
k = KubernetesPodOperator(task_id="task" +
self.get_current_task_name(),
pod_template_file=path,
do_xcom_push=True)
monitor_mock.return_value = (State.SUCCESS, None)
context = create_context(k)
with self.assertLogs(k.log, level=logging.DEBUG) as cm:
k.execute(context)
expected_line = textwrap.dedent("""\
DEBUG:airflow.task.operators:Starting pod:
api_version: v1
kind: Pod
metadata:
annotations: {}
cluster_name: null
creation_timestamp: null
deletion_grace_period_seconds: null\
""").strip()
self.assertTrue(
any(line.startswith(expected_line) for line in cm.output))
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
expected_dict = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {},
'labels': {},
'name': 'memory-demo',
'namespace': 'mem-example'
},
'spec': {
'affinity': {},
'containers': [
{
'args':
['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'env': [],
'envFrom': [],
'image':
'apache/airflow:stress-2020.07.10-1.0.4',
'name':
'base',
'ports': [],
'resources': {
'limits': {
'memory': '200Mi'
},
'requests': {
'memory': '100Mi'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}],
},
{
'command': [
'sh', '-c',
'trap "exit 0" INT; while true; do sleep 30; done;'
],
'image':
'alpine',
'name':
'airflow-xcom-sidecar',
'resources': {
'requests': {
'cpu': '1m'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}],
},
],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'restartPolicy':
'Never',
'securityContext': {},
'serviceAccountName':
'default',
'tolerations': [],
'volumes': [{
'emptyDir': {},
'name': 'xcom'
}],
},
}
self.assertEqual(expected_dict, actual_pod)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.start_pod")
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_priority_class_name(
self,
mock_client,
monitor_mock,
start_mock # pylint: disable=unused-argument
):
"""Test ability to assign priorityClassName to pod"""
from airflow.utils.state import State
priority_class_name = "medium-test"
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,
priority_class_name=priority_class_name,
)
monitor_mock.return_value = (State.SUCCESS, None)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['priorityClassName'] = priority_class_name
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_name(self):
pod_name_too_long = "a" * 221
with self.assertRaises(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,
)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
def test_on_kill(self, monitor_mock): # pylint: disable=unused-argument
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=["sleep 1000"],
labels={"foo": "bar"},
name="test",
task_id=name,
in_cluster=False,
do_xcom_push=False,
termination_grace_period=0,
)
context = create_context(k)
monitor_mock.return_value = (State.SUCCESS, None)
k.execute(context)
name = k.pod.metadata.name
pod = client.read_namespaced_pod(name=name, namespace=namespace)
self.assertEqual(pod.status.phase, "Running")
k.on_kill()
with self.assertRaises(ApiException):
pod = client.read_namespaced_pod(name=name, namespace=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.kubernetes.pod_launcher.PodLauncher.monitor_pod"
) as monitor_mock:
monitor_mock.return_value = (State.SUCCESS, 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 self.assertRaises(AirflowException):
k.execute(context)
pod = client.read_namespaced_pod(name=name, namespace=namespace)
self.assertEqual(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()
# Initialise client for the REST API used doing configuration.
#
# XXX Currently have a workaround here for OpenShift 4.0 beta versions
# which disables verification of the certificate. If don't use this the
# Python openshift/kubernetes clients will fail. We also disable any
# warnings from urllib3 to get rid of the noise in the logs this creates.
load_incluster_config()
import urllib3
urllib3.disable_warnings()
instance = Configuration()
instance.verify_ssl = False
Configuration.set_default(instance)
api_client = DynamicClient(ApiClient())
image_stream_resource = api_client.resources.get(
api_version='image.openshift.io/v1', kind='ImageStream')
# Helper function for determining the correct name for the image. We
# need to do this for references to image streams because of the image
# lookup policy often not being correctly setup on OpenShift clusters.
def resolve_image_name(name):
# If the image name contains a slash, we assume it is already
# referring to an image on some image registry. Even if it does
# not contain a slash, it may still be hosted on docker.io.
if name.find('/') != -1:
class TestKubernetesPodOperator(unittest.TestCase):
def setUp(self):
self.maxDiff = None # pylint: disable=invalid-name
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-')
}
},
'spec': {
'affinity': {},
'containers': [{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'imagePullPolicy': 'IfNotPresent',
'envFrom': [],
'name': 'base',
'ports': [],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy':
'Never',
'securityContext': {},
'serviceAccountName':
'default',
'tolerations': [],
'volumes': [],
}
}
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_config_path_move(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,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod, actual_pod)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
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,
)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_image_pull_secrets_correctly_set(self, mock_client,
launcher_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',
)
launcher_mock.return_value = (State.SUCCESS, None)
k.execute(None)
self.assertEqual(launcher_mock.call_args[0][0].spec.image_pull_secrets,
[k8s.V1LocalObjectReference(name=fake_pull_secrets)])
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.delete_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_delete_even_on_launcher_error(self, mock_client,
delete_pod_mock, run_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,
)
run_pod_mock.side_effect = AirflowException('fake failure')
with self.assertRaises(AirflowException):
k.execute(None)
assert delete_pod_mock.called
def test_working_pod(self):
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,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod, actual_pod)
def test_delete_operator_pod(self):
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,
is_delete_operator_pod=True,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_hostnetwork(self):
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,
hostnetwork=True,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['hostNetwork'] = True
self.assertEqual(self.expected_pod, actual_pod)
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",
task_id="task",
in_cluster=False,
do_xcom_push=False,
hostnetwork=True,
dnspolicy=dns_policy)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['hostNetwork'] = True
self.expected_pod['spec']['dnsPolicy'] = dns_policy
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_schedulername(self):
scheduler_name = "default-scheduler"
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,
schedulername=scheduler_name)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['schedulerName'] = scheduler_name
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_node_selectors(self):
node_selectors = {'beta.kubernetes.io/os': 'linux'}
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,
node_selectors=node_selectors,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['nodeSelector'] = node_selectors
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_resources(self):
resources = {
'limit_cpu': 0.25,
'limit_memory': '64Mi',
'request_cpu': '250m',
'request_memory': '64Mi',
}
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,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['resources'] = {
'requests': {
'memory': '64Mi',
'cpu': '250m'
},
'limits': {
'memory': '64Mi',
'cpu': 0.25,
'nvidia.com/gpu': None
}
}
self.assertEqual(self.expected_pod, actual_pod)
def test_pod_affinity(self):
affinity = {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': {
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'beta.kubernetes.io/os',
'operator': 'In',
'values': ['linux']
}]
}]
}
}
}
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,
affinity=affinity,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['affinity'] = affinity
self.assertEqual(self.expected_pod, actual_pod)
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],
)
k.execute(None)
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_volume_mount(self):
with mock.patch.object(PodLauncher, 'log') as mock_logger:
volume_mount = VolumeMount('test-volume',
mount_path='/root/mount_file',
sub_path=None,
read_only=True)
volume_config = {
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}
volume = Volume(name='test-volume', configs=volume_config)
args = ["cat /root/mount_file/test.txt"]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
volume_mounts=[volume_mount],
volumes=[volume],
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
k.execute(None)
mock_logger.info.assert_any_call(b"retrieved from mount\n")
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['args'] = args
self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{
'name':
'test-volume',
'mountPath':
'/root/mount_file',
'readOnly':
True
}]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
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",
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_run_as_user_non_root(self):
security_context = {
'securityContext': {
'runAsUser': 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_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_faulty_image(self):
bad_image_name = "foobar"
k = KubernetesPodOperator(
namespace='default',
image=bad_image_name,
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,
)
with self.assertRaises(AirflowException):
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0][
'image'] = bad_image_name
self.assertEqual(self.expected_pod, actual_pod)
def test_faulty_service_account(self):
bad_service_account_name = "foobar"
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,
startup_timeout_seconds=5,
service_account_name=bad_service_account_name,
)
with self.assertRaises(ApiException):
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec'][
'serviceAccountName'] = bad_service_account_name
self.assertEqual(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",
task_id="task",
in_cluster=False,
do_xcom_push=False,
)
with self.assertRaises(AirflowException):
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0][
'args'] = bad_internal_command
self.assertEqual(self.expected_pod, actual_pod)
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)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_envs_from_configmaps(self, mock_client, mock_launcher):
# GIVEN
from airflow.utils.state import State
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_launcher.return_value = (State.SUCCESS, None)
k.execute(None)
self.assertEqual(
mock_launcher.call_args[0][0].spec.containers[0].env_from, [
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name=configmap))
])
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_envs_from_secrets(self, mock_client, launcher_mock):
# GIVEN
from airflow.utils.state import State
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
launcher_mock.return_value = (State.SUCCESS, None)
k.execute(None)
self.assertEqual(
launcher_mock.call_args[0][0].spec.containers[0].env_from, [
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(
name=secret_ref))
])
def test_init_container(self):
# GIVEN
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=volume_mounts,
command=["bash", "-cx"],
args=["echo 10"])
volume_config = {'persistentVolumeClaim': {'claimName': 'test-volume'}}
volume = Volume(name='test-volume', configs=volume_config)
expected_init_container = {
'name':
'init-container',
'image':
'ubuntu:16.04',
'command': ['bash', '-cx'],
'args': ['echo 10'],
'env': [{
'name': 'key1',
'value': 'value1'
}, {
'name': 'key2',
'value': 'value2'
}],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'test-volume',
'readOnly': True
}],
}
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
volumes=[volume],
init_containers=[init_container],
in_cluster=False,
do_xcom_push=False,
)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['initContainers'] = [expected_init_container]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
self.assertEqual(self.expected_pod, actual_pod)
@mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.run_pod")
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_template_file(self, mock_client, launcher_mock):
from airflow.utils.state import State
k = KubernetesPodOperator(
task_id='task',
pod_template_file='tests/kubernetes/pod.yaml',
do_xcom_push=True)
launcher_mock.return_value = (State.SUCCESS, None)
k.execute(None)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': ANY,
'namespace': 'mem-example'
},
'spec': {
'volumes': [{
'name': 'xcom',
'emptyDir': {}
}],
'containers': [{
'args':
['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'image':
'polinux/stress',
'name':
'memory-demo-ctr',
'resources': {
'limits': {
'memory': '200Mi'
},
'requests': {
'memory': '100Mi'
}
},
'volumeMounts': [{
'name': 'xcom',
'mountPath': '/airflow/xcom'
}]
}, {
'name':
'airflow-xcom-sidecar',
'image':
"alpine",
'command': ['sh', '-c', PodDefaults.XCOM_CMD],
'volumeMounts': [{
'name': 'xcom',
'mountPath': '/airflow/xcom'
}],
'resources': {
'requests': {
'cpu': '1m'
}
},
}],
}
}, actual_pod)
class TestKubernetesPodOperatorSystem(unittest.TestCase):
def get_current_task_name(self):
# reverse test name to make pod name unique (it has limited length)
return "_" + unittest.TestCase.id(self).replace(".", "_")[::-1]
def setUp(self):
self.maxDiff = None
self.api_client = ApiClient()
self.expected_pod = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'default',
'name': ANY,
'annotations': {},
'labels': {
'foo': 'bar',
'kubernetes_pod_operator': 'True',
'airflow_version': airflow_version.replace('+', '-'),
'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b',
'dag_id': 'dag',
'task_id': ANY,
'try_number': '1',
},
},
'spec': {
'affinity': {},
'containers': [{
'image': 'ubuntu:16.04',
'args': ["echo 10"],
'command': ["bash", "-cx"],
'env': [],
'envFrom': [],
'resources': {},
'name': 'base',
'ports': [],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy':
'Never',
'securityContext': {},
'tolerations': [],
'volumes': [],
},
}
def tearDown(self) -> None:
client = kube_client.get_kube_client(in_cluster=False)
client.delete_collection_namespaced_pod(namespace="default")
import time
time.sleep(1)
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,
)
assert not k.do_xcom_push
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)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
assert self.expected_pod == actual_pod
def test_working_pod(self):
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,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
assert self.expected_pod['spec'] == actual_pod['spec']
assert self.expected_pod['metadata']['labels'] == actual_pod[
'metadata']['labels']
def test_delete_operator_pod(self):
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,
is_delete_operator_pod=True,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
assert self.expected_pod['spec'] == actual_pod['spec']
assert self.expected_pod['metadata']['labels'] == actual_pod[
'metadata']['labels']
def test_pod_hostnetwork(self):
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,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['hostNetwork'] = True
assert self.expected_pod['spec'] == actual_pod['spec']
assert self.expected_pod['metadata']['labels'] == actual_pod[
'metadata']['labels']
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_pod_schedulername(self):
scheduler_name = "default-scheduler"
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,
schedulername=scheduler_name,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['schedulerName'] = scheduler_name
assert self.expected_pod == actual_pod
def test_pod_node_selectors(self):
node_selectors = {'beta.kubernetes.io/os': 'linux'}
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,
node_selectors=node_selectors,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['nodeSelector'] = node_selectors
assert self.expected_pod == actual_pod
def test_pod_resources(self):
resources = k8s.V1ResourceRequirements(
requests={
'memory': '64Mi',
'cpu': '250m',
'ephemeral-storage': '1Gi'
},
limits={
'memory': '64Mi',
'cpu': 0.25,
'nvidia.com/gpu': None,
'ephemeral-storage': '2Gi'
},
)
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,
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,
'nvidia.com/gpu': None,
'ephemeral-storage': '2Gi'
},
}
assert self.expected_pod == actual_pod
def test_pod_affinity(self):
affinity = {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': {
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'beta.kubernetes.io/os',
'operator': 'In',
'values': ['linux']
}]
}]
}
}
}
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,
affinity=affinity,
)
context = create_context(k)
k.execute(context=context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['affinity'] = affinity
assert self.expected_pod == actual_pod
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_volume_mount(self):
with mock.patch.object(PodManager, 'log') as mock_logger:
volume_mount = k8s.V1VolumeMount(name='test-volume',
mount_path='/tmp/test_volume',
sub_path=None,
read_only=False)
volume = k8s.V1Volume(
name='test-volume',
persistent_volume_claim=k8s.
V1PersistentVolumeClaimVolumeSource(claim_name='test-volume'),
)
args = [
"echo \"retrieved from mount\" > /tmp/test_volume/test.txt "
"&& cat /tmp/test_volume/test.txt"
]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
volume_mounts=[volume_mount],
volumes=[volume],
name="test-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context=context)
mock_logger.info.assert_any_call('retrieved from mount')
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['containers'][0]['args'] = args
self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{
'name':
'test-volume',
'mountPath':
'/tmp/test_volume',
'readOnly':
False
}]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
assert 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_run_as_user_non_root(self):
security_context = {
'securityContext': {
'runAsUser': 1000,
}
}
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_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-fs-group",
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_faulty_image(self):
bad_image_name = "foobar"
k = KubernetesPodOperator(
namespace='default',
image=bad_image_name,
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,
startup_timeout_seconds=5,
)
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][
'image'] = bad_image_name
assert self.expected_pod == actual_pod
def test_faulty_service_account(self):
bad_service_account_name = "foobar"
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,
startup_timeout_seconds=5,
service_account_name=bad_service_account_name,
)
context = create_context(k)
pod = k.build_pod_request_obj(context)
with pytest.raises(
ApiException,
match="error looking up service account default/foobar"):
k.get_or_create_pod(pod, context)
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
@mock.patch("airflow.models.taskinstance.TaskInstance.xcom_push")
def test_xcom_push(self, xcom_push):
return_value = '{"foo": "bar"\n, "buzz": 2}'
args = [f'echo \'{return_value}\' > /airflow/xcom/return.json']
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=args,
labels={"foo": "bar"},
name="test-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=True,
)
context = create_context(k)
k.execute(context)
assert xcom_push.called_once_with(key=XCOM_RETURN_KEY,
value=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)
assert self.expected_pod == actual_pod
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_envs_from_secrets(self, mock_client, await_pod_completion_mock,
create_pod):
# 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-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
)
# THEN
await_pod_completion_mock.return_value = None
context = create_context(k)
with pytest.raises(AirflowException):
k.execute(context)
assert create_pod.call_args[1]['pod'].spec.containers[0].env_from == [
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(
name=secret_ref))
]
def test_env_vars(self):
# WHEN
env_vars = [
k8s.V1EnvVar(name="ENV1", value="val1"),
k8s.V1EnvVar(name="ENV2", value="val2"),
k8s.V1EnvVar(
name="ENV3",
value_from=k8s.V1EnvVarSource(
field_ref=k8s.V1ObjectFieldSelector(
field_path="status.podIP")),
),
]
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
env_vars=env_vars,
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,
)
# THEN
context = create_context(k)
actual_pod = self.api_client.sanitize_for_serialization(
k.build_pod_request_obj(context))
self.expected_pod['spec']['containers'][0]['env'] = [
{
'name': 'ENV1',
'value': 'val1'
},
{
'name': 'ENV2',
'value': 'val2'
},
{
'name': 'ENV3',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
},
]
assert self.expected_pod == actual_pod
def test_pod_template_file_system(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
assert result is not None
assert result == {"hello": "world"}
def test_pod_template_file_with_overrides_system(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
labels={
"foo": "bar",
"fizz": "buzz"
},
env_vars=[k8s.V1EnvVar(name="env_name", value="value")],
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
assert result is not None
assert k.pod.metadata.labels == {
'fizz': 'buzz',
'foo': 'bar',
'airflow_version': mock.ANY,
'dag_id': 'dag',
'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b',
'kubernetes_pod_operator': 'True',
'task_id': mock.ANY,
'try_number': '1',
}
assert k.pod.spec.containers[0].env == [
k8s.V1EnvVar(name="env_name", value="value")
]
assert result == {"hello": "world"}
def test_pod_template_file_with_full_pod_spec(self):
fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml'
pod_spec = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(labels={
"foo": "bar",
"fizz": "buzz"
}, ),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(
name="base",
env=[k8s.V1EnvVar(name="env_name", value="value")],
)
]),
)
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
pod_template_file=fixture,
full_pod_spec=pod_spec,
do_xcom_push=True,
)
context = create_context(k)
result = k.execute(context)
assert result is not None
assert k.pod.metadata.labels == {
'fizz': 'buzz',
'foo': 'bar',
'airflow_version': mock.ANY,
'dag_id': 'dag',
'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b',
'kubernetes_pod_operator': 'True',
'task_id': mock.ANY,
'try_number': '1',
}
assert k.pod.spec.containers[0].env == [
k8s.V1EnvVar(name="env_name", value="value")
]
assert result == {"hello": "world"}
def test_full_pod_spec(self):
pod_spec = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(labels={
"foo": "bar",
"fizz": "buzz"
},
namespace="default",
name="test-pod"),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
image="perl",
command=["/bin/bash"],
args=[
"-c",
'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json'
],
env=[k8s.V1EnvVar(name="env_name", value="value")],
)
],
restart_policy="Never",
),
)
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
in_cluster=False,
full_pod_spec=pod_spec,
do_xcom_push=True,
is_delete_operator_pod=False,
)
context = create_context(k)
result = k.execute(context)
assert result is not None
assert k.pod.metadata.labels == {
'fizz': 'buzz',
'foo': 'bar',
'airflow_version': mock.ANY,
'dag_id': 'dag',
'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b',
'kubernetes_pod_operator': 'True',
'task_id': mock.ANY,
'try_number': '1',
}
assert k.pod.spec.containers[0].env == [
k8s.V1EnvVar(name="env_name", value="value")
]
assert result == {"hello": "world"}
def test_init_container(self):
# GIVEN
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=volume_mounts,
command=["bash", "-cx"],
args=["echo 10"],
)
volume = k8s.V1Volume(
name='test-volume',
persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(
claim_name='test-volume'),
)
expected_init_container = {
'name':
'init-container',
'image':
'ubuntu:16.04',
'command': ['bash', '-cx'],
'args': ['echo 10'],
'env': [{
'name': 'key1',
'value': 'value1'
}, {
'name': 'key2',
'value': 'value2'
}],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'test-volume',
'readOnly': True
}],
}
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(),
volumes=[volume],
init_containers=[init_container],
in_cluster=False,
do_xcom_push=False,
)
context = create_context(k)
k.execute(context)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
self.expected_pod['spec']['initContainers'] = [expected_init_container]
self.expected_pod['spec']['volumes'] = [{
'name': 'test-volume',
'persistentVolumeClaim': {
'claimName': 'test-volume'
}
}]
assert self.expected_pod == actual_pod
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.extract_xcom"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_template_file(self, mock_client, await_pod_completion_mock,
create_mock, extract_xcom_mock):
extract_xcom_mock.return_value = '{}'
path = sys.path[0] + '/tests/kubernetes/pod.yaml'
k = KubernetesPodOperator(
task_id="task" + self.get_current_task_name(),
random_name_suffix=False,
pod_template_file=path,
do_xcom_push=True,
)
pod_mock = MagicMock()
pod_mock.status.phase = 'Succeeded'
await_pod_completion_mock.return_value = pod_mock
context = create_context(k)
with self.assertLogs(k.log, level=logging.DEBUG) as cm:
k.execute(context)
expected_line = textwrap.dedent("""\
DEBUG:airflow.task.operators:Starting pod:
api_version: v1
kind: Pod
metadata:
annotations: {}
cluster_name: null
creation_timestamp: null
deletion_grace_period_seconds: null\
""").strip()
assert any(line.startswith(expected_line) for line in cm.output)
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
expected_dict = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {},
'labels': {
'dag_id': 'dag',
'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b',
'kubernetes_pod_operator': 'True',
'task_id': mock.ANY,
'try_number': '1',
},
'name': 'memory-demo',
'namespace': 'mem-example',
},
'spec': {
'affinity': {},
'containers': [
{
'args':
['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'env': [],
'envFrom': [],
'image':
'ghcr.io/apache/airflow-stress:1.0.4-2021.07.04',
'name':
'base',
'ports': [],
'resources': {
'limits': {
'memory': '200Mi'
},
'requests': {
'memory': '100Mi'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}],
},
{
'command': [
'sh', '-c',
'trap "exit 0" INT; while true; do sleep 1; done;'
],
'image':
'alpine',
'name':
'airflow-xcom-sidecar',
'resources': {
'requests': {
'cpu': '1m'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}],
},
],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [],
'nodeSelector': {},
'restartPolicy':
'Never',
'securityContext': {},
'tolerations': [],
'volumes': [{
'emptyDir': {},
'name': 'xcom'
}],
},
}
version = actual_pod['metadata']['labels']['airflow_version']
assert version.startswith(airflow_version)
del actual_pod['metadata']['labels']['airflow_version']
assert expected_dict == actual_pod
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod"
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion"
)
@mock.patch("airflow.kubernetes.kube_client.get_kube_client")
def test_pod_priority_class_name(self, mock_client,
await_pod_completion_mock, create_mock):
"""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-" + str(random.randint(0, 1000000)),
task_id="task" + self.get_current_task_name(),
in_cluster=False,
do_xcom_push=False,
priority_class_name=priority_class_name,
)
pod_mock = MagicMock()
pod_mock.status.phase = 'Succeeded'
await_pod_completion_mock.return_value = pod_mock
context = 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_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,
)
@mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion"
)
def test_on_kill(self, await_pod_completion_mock):
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=["sleep 1000"],
labels={"foo": "bar"},
name="test",
task_id=name,
in_cluster=False,
do_xcom_push=False,
get_logs=False,
termination_grace_period=0,
)
context = create_context(k)
with pytest.raises(AirflowException):
k.execute(context)
name = k.pod.metadata.name
pod = client.read_namespaced_pod(name=name, namespace=namespace)
assert pod.status.phase == "Running"
k.on_kill()
with pytest.raises(ApiException,
match=r'pods \\"test.[a-z0-9]+\\" not found'):
client.read_namespaced_pod(name=name, namespace=namespace)
def test_reattach_failing_pod_once(self):
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)
# launch pod
with mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_pod_completion"
) as await_pod_completion_mock:
pod_mock = MagicMock()
# we don't want failure because we don't want the pod to be patched as "already_checked"
pod_mock.status.phase = 'Succeeded'
await_pod_completion_mock.return_value = pod_mock
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)
assert 'already_checked' not in pod.metadata.labels
# should not call `create_pod`, because there's a pod there it should find
# should use the found pod and patch as "already_checked" (in failure block)
with mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod"
) as create_mock:
with pytest.raises(AirflowException):
k.execute(context)
pod = client.read_namespaced_pod(name=name, namespace=namespace)
assert pod.metadata.labels["already_checked"] == "True"
create_mock.assert_not_called()
# `create_pod` should be called because though there's still a pod to be found,
# it will be `already_checked`
with mock.patch(
"airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.create_pod"
) as create_mock:
with pytest.raises(AirflowException):
k.execute(context)
create_mock.assert_called_once()
def __init__(self, _):
self.pod_dict = {}
self.pvc_list = []
self.api_client = ApiClient()
class TestKubernetesWorkerConfiguration(unittest.TestCase):
"""
Tests that if dags_volume_subpath/logs_volume_subpath configuration
options are passed to worker pod config
"""
affinity_config = {
'podAntiAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': [{
'topologyKey': 'kubernetes.io/hostname',
'labelSelector': {
'matchExpressions': [{
'key': 'app',
'operator': 'In',
'values': ['airflow']
}]
}
}]
}
}
tolerations_config = [{
'key': 'dedicated',
'operator': 'Equal',
'value': 'airflow'
}, {
'key': 'prod',
'operator': 'Exists'
}]
worker_annotations_config = {
'iam.amazonaws.com/role': 'role-arn',
'other/annotation': 'value'
}
def setUp(self):
if AirflowKubernetesScheduler is None:
self.skipTest("kubernetes python package is not installed")
self.kube_config = mock.MagicMock()
self.kube_config.airflow_home = '/'
self.kube_config.airflow_dags = 'dags'
self.kube_config.airflow_logs = 'logs'
self.kube_config.dags_volume_subpath = None
self.kube_config.logs_volume_subpath = None
self.kube_config.dags_in_image = False
self.kube_config.dags_folder = None
self.kube_config.git_dags_folder_mount_point = None
self.kube_config.kube_labels = {
'dag_id': 'original_dag_id',
'my_label': 'label_id'
}
self.api_client = ApiClient()
def test_worker_configuration_no_subpaths(self):
self.kube_config.dags_volume_claim = 'airflow-dags'
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
volume_mounts = worker_config._get_volume_mounts()
for volume_or_mount in volumes + volume_mounts:
if volume_or_mount.name != 'airflow-config':
self.assertNotIn(
'subPath',
self.api_client.sanitize_for_serialization(
volume_or_mount), "subPath shouldn't be defined")
@conf_vars({
('kubernetes', 'git_ssh_known_hosts_configmap_name'):
'airflow-configmap',
('kubernetes', 'git_ssh_key_secret_name'):
'airflow-secrets',
('kubernetes', 'git_user'):
'some-user',
('kubernetes', 'git_password'):
'some-password',
('kubernetes', 'git_repo'):
'[email protected]:apache/airflow.git',
('kubernetes', 'git_branch'):
'master',
('kubernetes', 'git_dags_folder_mount_point'):
'/usr/local/airflow/dags',
('kubernetes', 'delete_worker_pods'):
'True',
('kubernetes', 'kube_client_request_args'):
'{"_request_timeout" : [60,360]}',
})
def test_worker_configuration_auth_both_ssh_and_user(self):
with self.assertRaisesRegex(
AirflowConfigException,
'either `git_user` and `git_password`.*'
'or `git_ssh_key_secret_name`.*'
'but not both$'):
KubeConfig()
def test_worker_with_subpaths(self):
self.kube_config.dags_volume_subpath = 'dags'
self.kube_config.logs_volume_subpath = 'logs'
self.kube_config.dags_volume_claim = 'dags'
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
volume_mounts = worker_config._get_volume_mounts()
for volume in volumes:
self.assertNotIn(
'subPath', self.api_client.sanitize_for_serialization(volume),
"subPath isn't valid configuration for a volume")
for volume_mount in volume_mounts:
if volume_mount.name != 'airflow-config':
self.assertIn(
'subPath',
self.api_client.sanitize_for_serialization(volume_mount),
"subPath should've been passed to volumeMount configuration"
)
def test_worker_generate_dag_volume_mount_path(self):
self.kube_config.git_dags_folder_mount_point = '/root/airflow/git/dags'
self.kube_config.dags_folder = '/root/airflow/dags'
worker_config = WorkerConfiguration(self.kube_config)
self.kube_config.dags_volume_claim = 'airflow-dags'
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = '/host/airflow/dags'
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path,
self.kube_config.git_dags_folder_mount_point)
def test_worker_environment_no_dags_folder(self):
self.kube_config.airflow_configmap = ''
self.kube_config.git_dags_folder_mount_point = ''
self.kube_config.dags_folder = ''
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertNotIn('AIRFLOW__CORE__DAGS_FOLDER', env)
def test_worker_environment_when_dags_folder_specified(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_dags_folder_mount_point = ''
dags_folder = '/workers/path/to/dags'
self.kube_config.dags_folder = dags_folder
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_worker_environment_dags_folder_using_git_sync(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'dags'
self.kube_config.git_dags_folder_mount_point = '/workers/path/to/dags'
dags_folder = '{}/{}/{}'.format(
self.kube_config.git_dags_folder_mount_point,
self.kube_config.git_sync_dest, self.kube_config.git_subpath)
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_init_environment_using_git_sync_ssh_without_known_hosts(self):
# Tests the init environment created with git-sync SSH authentication option is correct
# without known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_secret_name = 'airflow-secrets'
self.kube_config.git_ssh_known_hosts_configmap_name = None
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0].env
self.assertIn(
k8s.V1EnvVar(name='GIT_SSH_KEY_FILE', value='/etc/git-secret/ssh'),
env)
self.assertIn(k8s.V1EnvVar(name='GIT_KNOWN_HOSTS', value='false'), env)
self.assertIn(k8s.V1EnvVar(name='GIT_SYNC_SSH', value='true'), env)
def test_init_environment_using_git_sync_ssh_with_known_hosts(self):
# Tests the init environment created with git-sync SSH authentication option is correct
# with known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
self.kube_config.git_ssh_known_hosts_configmap_name = 'airflow-configmap'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0].env
self.assertIn(
k8s.V1EnvVar(name='GIT_SSH_KEY_FILE', value='/etc/git-secret/ssh'),
env)
self.assertIn(k8s.V1EnvVar(name='GIT_KNOWN_HOSTS', value='true'), env)
self.assertIn(
k8s.V1EnvVar(name='GIT_SSH_KNOWN_HOSTS_FILE',
value='/etc/git-secret/known_hosts'), env)
self.assertIn(k8s.V1EnvVar(name='GIT_SYNC_SSH', value='true'), env)
def test_init_environment_using_git_sync_user_without_known_hosts(self):
# Tests the init environment created with git-sync User authentication option is correct
# without known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_user = '******'
self.kube_config.git_password = '******'
self.kube_config.git_ssh_known_hosts_configmap_name = None
self.kube_config.git_ssh_key_secret_name = None
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0].env
self.assertNotIn(
k8s.V1EnvVar(name='GIT_SSH_KEY_FILE', value='/etc/git-secret/ssh'),
env)
self.assertIn(k8s.V1EnvVar(name='GIT_SYNC_USERNAME', value='git_user'),
env)
self.assertIn(
k8s.V1EnvVar(name='GIT_SYNC_PASSWORD', value='git_password'), env)
self.assertIn(k8s.V1EnvVar(name='GIT_KNOWN_HOSTS', value='false'), env)
self.assertNotIn(
k8s.V1EnvVar(name='GIT_SSH_KNOWN_HOSTS_FILE',
value='/etc/git-secret/known_hosts'), env)
self.assertNotIn(k8s.V1EnvVar(name='GIT_SYNC_SSH', value='true'), env)
def test_init_environment_using_git_sync_user_with_known_hosts(self):
# Tests the init environment created with git-sync User authentication option is correct
# with known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_user = '******'
self.kube_config.git_password = '******'
self.kube_config.git_ssh_known_hosts_configmap_name = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = None
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0].env
self.assertNotIn(
k8s.V1EnvVar(name='GIT_SSH_KEY_FILE', value='/etc/git-secret/ssh'),
env)
self.assertIn(k8s.V1EnvVar(name='GIT_SYNC_USERNAME', value='git_user'),
env)
self.assertIn(
k8s.V1EnvVar(name='GIT_SYNC_PASSWORD', value='git_password'), env)
self.assertIn(k8s.V1EnvVar(name='GIT_KNOWN_HOSTS', value='true'), env)
self.assertIn(
k8s.V1EnvVar(name='GIT_SSH_KNOWN_HOSTS_FILE',
value='/etc/git-secret/known_hosts'), env)
self.assertNotIn(k8s.V1EnvVar(name='GIT_SYNC_SSH', value='true'), env)
def test_init_environment_using_git_sync_run_as_user_empty(self):
# Tests if git_syn_run_as_user is none, then no securityContext created in init container
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.git_sync_run_as_user = ''
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
self.assertIsNone(init_containers[0].security_context)
def test_init_environment_using_git_sync_run_as_user_root(self):
# Tests if git_syn_run_as_user is '0', securityContext is created with
# the right uid
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.git_sync_run_as_user = 0
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
self.assertEqual(0, init_containers[0].security_context.run_as_user)
def test_make_pod_run_as_user_0(self):
# Tests the pod created with run-as-user 0 actually gets that in it's config
self.kube_config.worker_run_as_user = 0
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
self.assertEqual(0, pod.spec.security_context.run_as_user)
def test_make_pod_assert_labels(self):
# Tests the pod created has all the expected labels set
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", "sample-uuid", "test_pod_id",
"test_dag_id", "test_task_id",
"2019-11-21 11:08:22.920875", 1,
"bash -c 'ls /'")
expected_labels = {
'airflow-worker': 'sample-uuid',
'airflow_version': airflow_version.replace('+', '-'),
'dag_id': 'test_dag_id',
'execution_date': '2019-11-21 11:08:22.920875',
'kubernetes_executor': 'True',
'task_id': 'test_task_id',
'try_number': '1'
}
self.assertEqual(pod.metadata.labels, expected_labels)
def test_make_pod_git_sync_ssh_without_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct without known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
init_containers = worker_config._get_init_containers()
git_ssh_key_file = next(
(x.value
for x in init_containers[0].env if x.name == 'GIT_SSH_KEY_FILE'),
None)
volume_mount_ssh_key = next(
(x.mount_path for x in init_containers[0].volume_mounts
if x.name == worker_config.git_sync_ssh_secret_volume_name), None)
self.assertTrue(git_ssh_key_file)
self.assertTrue(volume_mount_ssh_key)
self.assertEqual(65533, pod.spec.security_context.fs_group)
self.assertEqual(
git_ssh_key_file, volume_mount_ssh_key,
'The location where the git ssh secret is mounted'
' needs to be the same as the GIT_SSH_KEY_FILE path')
def test_make_pod_git_sync_credentials_secret(self):
# Tests the pod created with git_sync_credentials_secret will get into the init container
self.kube_config.git_sync_credentials_secret = 'airflow-git-creds-secret'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
username_env = k8s.V1EnvVar(
name='GIT_SYNC_USERNAME',
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
name=self.kube_config.git_sync_credentials_secret,
key='GIT_SYNC_USERNAME')))
password_env = k8s.V1EnvVar(
name='GIT_SYNC_PASSWORD',
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
name=self.kube_config.git_sync_credentials_secret,
key='GIT_SYNC_PASSWORD')))
self.assertIn(
username_env, pod.spec.init_containers[0].env,
'The username env for git credentials did not get into the init container'
)
self.assertIn(
password_env, pod.spec.init_containers[0].env,
'The password env for git credentials did not get into the init container'
)
def test_make_pod_git_sync_rev(self):
# Tests the pod created with git_sync_credentials_secret will get into the init container
self.kube_config.git_sync_rev = 'sampletag'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
rev_env = k8s.V1EnvVar(
name='GIT_SYNC_REV',
value=self.kube_config.git_sync_rev,
)
self.assertIn(
rev_env, pod.spec.init_containers[0].env,
'The git_sync_rev env did not get into the init container')
def test_make_pod_git_sync_ssh_with_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct with known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
git_ssh_known_hosts_file = next(
(x.value for x in init_containers[0].env
if x.name == 'GIT_SSH_KNOWN_HOSTS_FILE'), None)
volume_mount_ssh_known_hosts_file = next(
(x.mount_path for x in init_containers[0].volume_mounts
if x.name == worker_config.git_sync_ssh_known_hosts_volume_name),
None)
self.assertTrue(git_ssh_known_hosts_file)
self.assertTrue(volume_mount_ssh_known_hosts_file)
self.assertEqual(
git_ssh_known_hosts_file, volume_mount_ssh_known_hosts_file,
'The location where the git known hosts file is mounted'
' needs to be the same as the GIT_SSH_KNOWN_HOSTS_FILE path')
def test_make_pod_with_empty_executor_config(self):
self.kube_config.kube_affinity = self.affinity_config
self.kube_config.kube_tolerations = self.tolerations_config
self.kube_config.kube_annotations = self.worker_annotations_config
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
self.assertTrue(pod.spec.affinity['podAntiAffinity'] is not None)
self.assertEqual(
'app', pod.spec.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']['matchExpressions'][0]['key'])
self.assertEqual(2, len(pod.spec.tolerations))
self.assertEqual('prod', pod.spec.tolerations[1]['key'])
self.assertEqual('role-arn',
pod.metadata.annotations['iam.amazonaws.com/role'])
self.assertEqual('value', pod.metadata.annotations['other/annotation'])
def test_make_pod_with_executor_config(self):
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
config_pod = PodGenerator(
image='',
affinity=self.affinity_config,
tolerations=self.tolerations_config,
).gen_pod()
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
result = PodGenerator.reconcile_pods(pod, config_pod)
self.assertTrue(result.spec.affinity['podAntiAffinity'] is not None)
self.assertEqual(
'app', result.spec.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']['matchExpressions'][0]['key'])
self.assertEqual(2, len(result.spec.tolerations))
self.assertEqual('prod', result.spec.tolerations[1]['key'])
def test_worker_pvc_dags(self):
# Tests persistence volume config created when `dags_volume_claim` is set
self.kube_config.dags_volume_claim = 'airflow-dags'
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
volume_mounts = worker_config._get_volume_mounts()
init_containers = worker_config._get_init_containers()
dag_volume = [
volume for volume in volumes if volume.name == 'airflow-dags'
]
dag_volume_mount = [
mount for mount in volume_mounts if mount.name == 'airflow-dags'
]
self.assertEqual('airflow-dags',
dag_volume[0].persistent_volume_claim.claim_name)
self.assertEqual(1, len(dag_volume_mount))
self.assertTrue(dag_volume_mount[0].read_only)
self.assertEqual(0, len(init_containers))
def test_worker_git_dags(self):
# Tests persistence volume config created when `git_repo` is set
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_folder = '/usr/local/airflow/dags'
self.kube_config.worker_dags_folder = '/usr/local/airflow/dags'
self.kube_config.git_sync_container_repository = 'gcr.io/google-containers/git-sync-amd64'
self.kube_config.git_sync_container_tag = 'v2.0.5'
self.kube_config.git_sync_container = 'gcr.io/google-containers/git-sync-amd64:v2.0.5'
self.kube_config.git_sync_init_container_name = 'git-sync-clone'
self.kube_config.git_subpath = 'dags_folder'
self.kube_config.git_sync_root = '/git'
self.kube_config.git_sync_run_as_user = 65533
self.kube_config.git_dags_folder_mount_point = '/usr/local/airflow/dags/repo/dags_folder'
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
volume_mounts = worker_config._get_volume_mounts()
dag_volume = [
volume for volume in volumes if volume.name == 'airflow-dags'
]
dag_volume_mount = [
mount for mount in volume_mounts if mount.name == 'airflow-dags'
]
self.assertIsNotNone(dag_volume[0].empty_dir)
self.assertEqual(self.kube_config.git_dags_folder_mount_point,
dag_volume_mount[0].mount_path)
self.assertTrue(dag_volume_mount[0].read_only)
init_container = worker_config._get_init_containers()[0]
init_container_volume_mount = [
mount for mount in init_container.volume_mounts
if mount.name == 'airflow-dags'
]
self.assertEqual('git-sync-clone', init_container.name)
self.assertEqual('gcr.io/google-containers/git-sync-amd64:v2.0.5',
init_container.image)
self.assertEqual(1, len(init_container_volume_mount))
self.assertFalse(init_container_volume_mount[0].read_only)
self.assertEqual(65533, init_container.security_context.run_as_user)
def test_worker_container_dags(self):
# Tests that the 'airflow-dags' persistence volume is NOT created when `dags_in_image` is set
self.kube_config.dags_in_image = True
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
volume_mounts = worker_config._get_volume_mounts()
dag_volume = [
volume for volume in volumes if volume.name == 'airflow-dags'
]
dag_volume_mount = [
mount for mount in volume_mounts if mount.name == 'airflow-dags'
]
init_containers = worker_config._get_init_containers()
self.assertEqual(0, len(dag_volume))
self.assertEqual(0, len(dag_volume_mount))
self.assertEqual(0, len(init_containers))
def test_set_airflow_local_settings_configmap(self):
"""
Test that airflow_local_settings.py can be set via configmap by
checking volume & volume-mounts are set correctly.
"""
self.kube_config.airflow_home = '/usr/local/airflow'
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.airflow_local_settings_configmap = 'airflow-configmap'
self.kube_config.dags_folder = '/workers/path/to/dags'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
airflow_config_volume = [
volume for volume in pod.spec.volumes
if volume.name == 'airflow-config'
]
# Test that volume_name is found
self.assertEqual(1, len(airflow_config_volume))
# Test that config map exists
self.assertEqual("airflow-configmap",
airflow_config_volume[0].config_map.name)
# Test Volume Mount exists
local_setting_volume_mount = [
volume_mount
for volume_mount in pod.spec.containers[0].volume_mounts
if volume_mount.name == 'airflow-config'
]
self.assertEqual(1, len(local_setting_volume_mount))
# Test Mounth Path is set correctly.
self.assertEqual('/usr/local/airflow/config/airflow_local_settings.py',
local_setting_volume_mount[0].mount_path)
self.assertEqual(True, local_setting_volume_mount[0].read_only)
self.assertEqual('airflow_local_settings.py',
local_setting_volume_mount[0].sub_path)
def test_kubernetes_environment_variables(self):
# Tests the kubernetes environment variables get copied into the worker pods
input_environment = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
for key in input_environment:
self.assertIn(key, env)
self.assertIn(input_environment[key], env.values())
core_executor = 'AIRFLOW__CORE__EXECUTOR'
input_environment = {core_executor: 'NotLocalExecutor'}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(env[core_executor], 'LocalExecutor')
def test_get_secrets(self):
# Test when secretRef is None and kube_secrets is not empty
self.kube_config.kube_secrets = {
'AWS_SECRET_KEY': 'airflow-secret=aws_secret_key',
'POSTGRES_PASSWORD': '******'
}
self.kube_config.env_from_secret_ref = None
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
secrets.sort(key=lambda secret: secret.deploy_target)
expected = [
Secret('env', 'AWS_SECRET_KEY', 'airflow-secret',
'aws_secret_key'),
Secret('env', 'POSTGRES_PASSWORD', 'airflow-secret',
'postgres_credentials')
]
self.assertListEqual(expected, secrets)
# Test when secret is not empty and kube_secrets is empty dict
self.kube_config.kube_secrets = {}
self.kube_config.env_from_secret_ref = 'secret_a,secret_b'
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
expected = [
Secret('env', None, 'secret_a'),
Secret('env', None, 'secret_b')
]
self.assertListEqual(expected, secrets)
def test_get_env_from(self):
# Test when configmap is empty
self.kube_config.env_from_configmap_ref = ''
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_env_from()
self.assertListEqual([], configmaps)
# test when configmap is not empty
self.kube_config.env_from_configmap_ref = 'configmap_a,configmap_b'
self.kube_config.env_from_secret_ref = 'secretref_a,secretref_b'
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_env_from()
self.assertListEqual([
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='secretref_a')),
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(
name='secretref_b'))
], configmaps)
def test_get_labels(self):
worker_config = WorkerConfiguration(self.kube_config)
labels = worker_config._get_labels(
{'my_kube_executor_label': 'kubernetes'}, {
'dag_id': 'override_dag_id',
})
self.assertEqual(
{
'my_label': 'label_id',
'dag_id': 'override_dag_id',
'my_kube_executor_label': 'kubernetes',
}, labels)
def test_make_pod_with_image_pull_secrets(self):
# Tests the pod created with image_pull_secrets actually gets that in it's config
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
self.kube_config.image_pull_secrets = 'image_pull_secret1,image_pull_secret2'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'")
self.assertEqual(2, len(pod.spec.image_pull_secrets))
