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_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.as_pod()
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_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.as_pod()
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_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_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_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_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.assertFalse({
'name': 'GIT_SSH_KEY_FILE',
'value': '/etc/git-secret/ssh'
} in env)
self.assertTrue({
'name': 'GIT_SYNC_USERNAME',
'value': 'git_user'
} in env)
self.assertTrue({
'name': 'GIT_SYNC_PASSWORD',
'value': 'git_password'
} in env)
self.assertTrue({'name': 'GIT_KNOWN_HOSTS', 'value': 'true'} in env)
self.assertTrue({
'name': 'GIT_SSH_KNOWN_HOSTS_FILE',
'value': '/etc/git-secret/known_hosts'
} in env)
self.assertFalse({'name': 'GIT_SYNC_SSH', 'value': 'true'} in env)
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
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[]
)
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 /'",
kube_executor_config)
self.assertTrue(pod.affinity['podAntiAffinity'] is not None)
self.assertEqual('app',
pod.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']
['matchExpressions'][0]
['key'])
self.assertEqual(2, len(pod.tolerations))
self.assertEqual('prod', pod.tolerations[1]['key'])
def test_set_airflow_configmap_different_for_local_setting(self):
"""
Test that airflow_local_settings.py can be set via configmap by
checking volume & volume-mounts are set correctly when using a different
configmap than airflow_configmap (airflow.cfg)
"""
self.kube_config.airflow_home = '/usr/local/airflow'
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.airflow_local_settings_configmap = 'airflow-ls-configmap'
self.kube_config.dags_folder = '/workers/path/to/dags'
worker_config = WorkerConfiguration(self.kube_config)
pod = worker_config.as_pod()
pod_spec_dict = pod.spec.to_dict()
airflow_local_settings_volume = [
volume for volume in pod_spec_dict['volumes'] if volume["name"] == 'airflow-local-settings'
]
# Test that volume_name is found
self.assertEqual(1, len(airflow_local_settings_volume))
# Test that config map exists
self.assertEqual(
{'default_mode': None, 'items': None, 'name': 'airflow-ls-configmap', 'optional': None},
airflow_local_settings_volume[0]['config_map']
)
# Test that 2 Volume Mounts exists and has 2 different mount-paths
# One for airflow.cfg
# Second for airflow_local_settings.py
airflow_cfg_volume_mount = [
volume_mount for volume_mount in pod_spec_dict['containers'][0]['volume_mounts']
if volume_mount['name'] == 'airflow-config'
]
local_setting_volume_mount = [
volume_mount for volume_mount in pod_spec_dict['containers'][0]['volume_mounts']
if volume_mount['name'] == 'airflow-local-settings'
]
self.assertEqual(1, len(airflow_cfg_volume_mount))
self.assertEqual(1, len(local_setting_volume_mount))
self.assertEqual(
[
{
'mount_path': '/usr/local/airflow/config/airflow_local_settings.py',
'mount_propagation': None,
'name': 'airflow-local-settings',
'read_only': True,
'sub_path': 'airflow_local_settings.py',
'sub_path_expr': None
}
],
local_setting_volume_mount
)
self.assertEqual(
[
{
'mount_path': '/usr/local/airflow/airflow.cfg',
'mount_propagation': None,
'name': 'airflow-config',
'read_only': True,
'sub_path': 'airflow.cfg',
'sub_path_expr': None
}
],
airflow_cfg_volume_mount
)
def test_worker_adds_config(self):
worker_config = WorkerConfiguration(self.kube_config)
volumes = worker_config._get_volumes()
print(volumes)
