def test_should_create_pgbouncer_service(self):
docs = render_chart(
values={"pgbouncer": {"enabled": True}},
show_only=["templates/pgbouncer/pgbouncer-service.yaml"],
)
assert "Service" == jmespath.search("kind", docs[0])
assert "RELEASE-NAME-pgbouncer" == jmespath.search("metadata.name", docs[0])
assert "true" == jmespath.search('metadata.annotations."prometheus.io/scrape"', docs[0])
assert "9127" == jmespath.search('metadata.annotations."prometheus.io/port"', docs[0])
assert {"prometheus.io/scrape": "true", "prometheus.io/port": "9127"} == jmespath.search(
"metadata.annotations", docs[0]
)
assert {"name": "pgbouncer", "protocol": "TCP", "port": 6543} in jmespath.search(
"spec.ports", docs[0]
)
assert {"name": "pgbouncer-metrics", "protocol": "TCP", "port": 9127} in jmespath.search(
"spec.ports", docs[0]
)
def test_should_set_the_dags_volume_claim_correctly_when_using_an_existing_claim(
self):
docs = render_chart(
values={
"dags": {
"persistence": {
"enabled": True,
"existingClaim": "test-claim"
}
}
},
show_only=["templates/pod-template-file.yaml"],
chart_dir=self.temp_chart_dir,
)
assert {
"name": "dags",
"persistentVolumeClaim": {
"claimName": "test-claim"
}
} in jmespath.search("spec.volumes", docs[0])
def test_should_add_env(self):
docs = render_chart(
values={
"airflowVersion": "1.10.14",
"dags": {
"gitSync": {
"enabled": True,
"env": [{
"name": "FOO",
"value": "bar"
}],
}
},
},
show_only=["templates/webserver/webserver-deployment.yaml"],
)
assert {
"name": "FOO",
"value": "bar"
} in jmespath.search("spec.template.spec.containers[1].env", docs[0])
def test_should_add_dags_volume_to_the_worker_if_git_sync_is_enabled_and_persistence_is_disabled(
self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"dags": {
"gitSync": {
"enabled": True
},
"persistence": {
"enabled": False
}
},
},
show_only=["templates/workers/worker-deployment.yaml"],
)
assert "config" == jmespath.search(
"spec.template.spec.volumes[0].name", docs[0])
assert "dags" == jmespath.search("spec.template.spec.volumes[1].name",
docs[0])
def test_logs_persistence_changes_volume(self, log_persistence_values,
expected_volume):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"workers": {
"persistence": {
"enabled": False
}
},
"logs": {
"persistence": log_persistence_values
},
},
show_only=["templates/workers/worker-deployment.yaml"],
)
assert {
"name": "logs",
**expected_volume
} == jmespath.search("spec.template.spec.volumes[1]", docs[0])
def test_validate_if_ssh_known_hosts_are_added(self):
docs = render_chart(
values={
"dags": {
"gitSync": {
"enabled": True,
"containerName": "git-sync-test",
"sshKeySecret": "ssh-secret",
"knownHosts": "github.com ssh-rsa AAAABdummy",
"branch": "test-branch",
}
}
},
show_only=["templates/pod-template-file.yaml"],
)
self.assertIn(
{
"name": "GIT_KNOWN_HOSTS",
"value": "true"
},
jmespath.search("spec.initContainers[0].env", docs[0]),
)
self.assertIn(
{
"name": "git-sync-known-hosts",
"configMap": {
"defaultMode": 288,
"name": "RELEASE-NAME-airflow-config"
},
},
jmespath.search("spec.volumes", docs[0]),
)
self.assertIn(
{
"name": "git-sync-known-hosts",
"mountPath": "/etc/git-secret/known_hosts",
"subPath": "known_hosts",
},
jmespath.search("spec.containers[0].volumeMounts", docs[0]),
)
def test_should_set_username_and_pass_env_variables(self):
docs = render_chart(
values={
"dags": {
"gitSync": {
"enabled": True,
"credentialsSecret": "user-pass-secret",
"sshKeySecret": None,
}
}
},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
assert {
"name": "GIT_SYNC_USERNAME",
"valueFrom": {"secretKeyRef": {"name": "user-pass-secret", "key": "GIT_SYNC_USERNAME"}},
} in jmespath.search("spec.template.spec.containers[1].env", docs[0])
assert {
"name": "GIT_SYNC_PASSWORD",
"valueFrom": {"secretKeyRef": {"name": "user-pass-secret", "key": "GIT_SYNC_PASSWORD"}},
} in jmespath.search("spec.template.spec.containers[1].env", docs[0])
def test_should_create_a_worker_deployment_with_the_celery_executo(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"dags": {
"persistence": {
"enabled": True
},
"gitSync": {
"enabled": True
}
},
},
show_only=["templates/workers/worker-deployment.yaml"],
)
self.assertEqual(
"config",
jmespath.search("spec.template.spec.volumes[0].name", docs[0]))
self.assertEqual(
"dags",
jmespath.search("spec.template.spec.volumes[1].name", docs[0]))
def test_should_create_flower_deployment_with_authorization(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"flower": {"username": "******", "password": "******"},
"ports": {"flowerUI": 7777},
},
show_only=["templates/flower/flower-deployment.yaml"],
)
self.assertEqual(
"AIRFLOW__CELERY__FLOWER_BASIC_AUTH",
jmespath.search("spec.template.spec.containers[0].env[0].name", docs[0]),
)
self.assertEqual(
['curl', '--user', '$AIRFLOW__CELERY__FLOWER_BASIC_AUTH', 'localhost:7777'],
jmespath.search("spec.template.spec.containers[0].livenessProbe.exec.command", docs[0]),
)
self.assertEqual(
['curl', '--user', '$AIRFLOW__CELERY__FLOWER_BASIC_AUTH', 'localhost:7777'],
jmespath.search("spec.template.spec.containers[0].readinessProbe.exec.command", docs[0]),
)
def test_should_add_extra_volume_and_extra_volume_mount(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"scheduler": {
"extraVolumes": [{
"name": "test-volume",
"emptyDir": {}
}],
"extraVolumeMounts": [{
"name": "test-volume",
"mountPath": "/opt/test"
}],
},
},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
assert "test-volume" == jmespath.search(
"spec.template.spec.volumes[1].name", docs[0])
assert "test-volume" == jmespath.search(
"spec.template.spec.containers[0].volumeMounts[3].name", docs[0])
def setUpClass(cls) -> None:
values_str = textwrap.dedent(
"""
executor: "CeleryExecutor"
extraEnvFrom: |
- secretRef:
name: '{{ .Release.Name }}-airflow-connections'
- configMapRef:
name: '{{ .Release.Name }}-airflow-variables'
extraEnv: |
- name: PLATFORM
value: FR
- name: TEST
valueFrom:
secretKeyRef:
name: '{{ .Release.Name }}-some-secret'
key: connection
"""
)
values = yaml.safe_load(values_str)
cls.k8s_objects = render_chart(RELEASE_NAME, values=values)
cls.k8s_objects_by_key = prepare_k8s_lookup_dict(cls.k8s_objects)
def test_workers_update_strategy(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"workers": {
"updateStrategy": {
"strategy": {
"rollingUpdate": {
"maxSurge": "100%",
"maxUnavailable": "50%"
}
}
},
},
},
show_only=["templates/workers/worker-deployment.yaml"],
)
assert "100%" == jmespath.search(
"spec.strategy.rollingUpdate.maxSurge", docs[0])
assert "50%" == jmespath.search(
"spec.strategy.rollingUpdate.maxUnavailable", docs[0])
def test_extra_configmaps(self):
values_str = textwrap.dedent("""
extraConfigMaps:
"{{ .Release.Name }}-airflow-variables":
data: |
AIRFLOW_VAR_HELLO_MESSAGE: "Hi!"
AIRFLOW_VAR_KUBERNETES_NAMESPACE: "{{ .Release.Namespace }}"
"{{ .Release.Name }}-other-variables":
data: |
HELLO_WORLD: "Hi again!"
""")
values = yaml.safe_load(values_str)
k8s_objects = render_chart(
RELEASE_NAME,
values=values,
show_only=["templates/configmaps/extra-configmaps.yaml"])
k8s_objects_by_key = prepare_k8s_lookup_dict(k8s_objects)
all_expected_keys = [
("ConfigMap", f"{RELEASE_NAME}-airflow-variables"),
("ConfigMap", f"{RELEASE_NAME}-other-variables"),
]
self.assertEqual(set(k8s_objects_by_key.keys()),
set(all_expected_keys))
all_expected_data = [
{
"AIRFLOW_VAR_HELLO_MESSAGE": "Hi!",
"AIRFLOW_VAR_KUBERNETES_NAMESPACE": "default"
},
{
"HELLO_WORLD": "Hi again!"
},
]
for expected_key, expected_data in zip(all_expected_keys,
all_expected_data):
configmap_obj = k8s_objects_by_key[expected_key]
self.assertEqual(configmap_obj["data"], expected_data)
def test_deployments_no_rbac_no_sa(self):
k8s_objects = render_chart(
"TEST-RBAC",
values={
"rbac": {
"create": False
},
"cleanup": {
"enabled": True,
"serviceAccount": {
"create": False,
},
},
"scheduler": {
"serviceAccount": {
"create": False
}
},
"webserver": {
"serviceAccount": {
"create": False
}
},
"workers": {
"serviceAccount": {
"create": False
}
},
},
)
list_of_kind_names_tuples = [(k8s_object['kind'],
k8s_object['metadata']['name'])
for k8s_object in k8s_objects]
self.assertCountEqual(
list_of_kind_names_tuples,
CLEANUP_DEPLOYMENT_KIND_NAME_TUPLES,
)
def test_service_account_without_resource(self):
k8s_objects = render_chart(
"TEST-RBAC",
values={
"fullnameOverride": "TEST-RBAC",
"executor": "LocalExecutor",
"cleanup": {
"enabled": False
},
"pgbouncer": {
"enabled": False
},
"redis": {
"enabled": False
},
"flower": {
"enabled": False
},
"statsd": {
"enabled": False
},
"webserver": {
"defaultUser": {
"enabled": False
}
},
},
)
list_of_sa_names = [
k8s_object['metadata']['name'] for k8s_object in k8s_objects
if k8s_object['kind'] == "ServiceAccount"
]
service_account_names = [
'TEST-RBAC-scheduler',
'TEST-RBAC-webserver',
'TEST-RBAC-migrate-database-job',
]
self.assertCountEqual(list_of_sa_names, service_account_names)
def test_extra_volume_and_git_sync_extra_volume_mount(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"scheduler": {
"extraVolumes": [{
"name": "test-volume",
"emptyDir": {}
}],
},
"dags": {
"gitSync": {
"enabled":
True,
"extraVolumeMounts": [{
"mountPath": "/opt/test",
"name": "test-volume"
}],
}
},
},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
assert {
"name": "test-volume",
"emptyDir": {}
} in jmespath.search("spec.template.spec.volumes", docs[0])
assert {
'mountPath': '/git',
'name': 'dags'
} in jmespath.search("spec.template.spec.containers[1].volumeMounts",
docs[0])
assert {
"name": "test-volume",
"mountPath": "/opt/test"
} in jmespath.search("spec.template.spec.containers[1].volumeMounts",
docs[0])
def test_should_set_username_and_pass_env_variables(self):
docs = render_chart(
values={
"dags": {
"gitSync": {
"enabled": True,
"credentialsSecret": "user-pass-secret",
"sshKeySecret": None,
}
}
},
show_only=["templates/pod-template-file.yaml"],
)
self.assertIn(
{
"name": "GIT_SYNC_USERNAME",
"valueFrom": {
"secretKeyRef": {
"name": "user-pass-secret",
"key": "GIT_SYNC_USERNAME"
}
},
},
jmespath.search("spec.initContainers[0].env", docs[0]),
)
self.assertIn(
{
"name": "GIT_SYNC_PASSWORD",
"valueFrom": {
"secretKeyRef": {
"name": "user-pass-secret",
"key": "GIT_SYNC_PASSWORD"
}
},
},
jmespath.search("spec.initContainers[0].env", docs[0]),
)
def test_redis_by_chart_password(self, executor):
k8s_objects = render_chart(
RELEASE_NAME_REDIS,
{
"executor": executor,
"networkPolicies": {"enabled": True},
"redis": {"enabled": True, "password": "******"},
},
)
k8s_obj_by_key = prepare_k8s_lookup_dict(k8s_objects)
created_redis_objects = SET_POSSIBLE_REDIS_OBJECT_KEYS & set(k8s_obj_by_key.keys())
assert created_redis_objects == SET_POSSIBLE_REDIS_OBJECT_KEYS
self.assert_password_and_broker_url_secrets(
k8s_obj_by_key,
expected_password_match="test-redis-password",
expected_broker_url_match=re.escape(
"redis://:test-redis-password%21%40%23$%25%5E&%2A%28%29_+@TEST-REDIS-redis:6379/0"
),
)
self.assert_broker_url_env(k8s_obj_by_key)
def test_should_add_volume_and_volume_mount_when_exist_webserver_config(
self):
docs = render_chart(
values={"webserver": {
"webserverConfig": "CSRF_ENABLED = True"
}},
show_only=["templates/webserver/webserver-deployment.yaml"],
)
assert {
"name": "webserver-config",
"configMap": {
"name": "RELEASE-NAME-webserver-config"
},
} in jmespath.search("spec.template.spec.volumes", docs[0])
assert {
"name": "webserver-config",
"mountPath": "/opt/airflow/webserver_config.py",
"subPath": "webserver_config.py",
"readOnly": True,
} in jmespath.search("spec.template.spec.containers[0].volumeMounts",
docs[0])
def test_deployments_with_rbac_with_sa(self):
k8s_objects = render_chart(
"TEST-RBAC",
values={
"fullnameOverride": "TEST-RBAC",
"cleanup": {
"enabled": True
},
"pgbouncer": {
"enabled": True
},
},
)
list_of_kind_names_tuples = [(k8s_object['kind'],
k8s_object['metadata']['name'])
for k8s_object in k8s_objects]
real_list_of_kind_names = (DEPLOYMENT_NO_RBAC_NO_SA_KIND_NAME_TUPLES +
SERVICE_ACCOUNT_NAME_TUPLES +
RBAC_ENABLED_KIND_NAME_TUPLES)
self.assertCountEqual(
list_of_kind_names_tuples,
real_list_of_kind_names,
)
def test_should_create_cronjob_for_enabled_cleanup(self):
docs = render_chart(
values={
"cleanup": {"enabled": True},
},
show_only=["templates/cleanup/cleanup-cronjob.yaml"],
)
assert "airflow-cleanup-pods" == jmespath.search(
"spec.jobTemplate.spec.template.spec.containers[0].name", docs[0]
)
assert "apache/airflow:2.0.0" == jmespath.search(
"spec.jobTemplate.spec.template.spec.containers[0].image", docs[0]
)
assert {"name": "config", "configMap": {"name": "RELEASE-NAME-airflow-config"}} in jmespath.search(
"spec.jobTemplate.spec.template.spec.volumes", docs[0]
)
assert {
"name": "config",
"mountPath": "/opt/airflow/airflow.cfg",
"subPath": "airflow.cfg",
"readOnly": True,
} in jmespath.search("spec.jobTemplate.spec.template.spec.containers[0].volumeMounts", docs[0])
def test_should_set_the_volume_claim_correctly_when_using_an_existing_claim(
self):
docs = render_chart(
values={
"dags": {
"persistence": {
"enabled": True,
"existingClaim": "test-claim"
}
}
},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
self.assertIn(
{
"name": "dags",
"persistentVolumeClaim": {
"claimName": "test-claim"
}
},
jmespath.search("spec.template.spec.volumes", docs[0]),
)
def test_keda_enabled(self, executor, is_created):
"""
ScaledObject should only be created when set to enabled and executor is Celery or CeleryKubernetes
"""
docs = render_chart(
values={
"workers": {
"keda": {
"enabled": True
},
"persistence": {
"enabled": False
}
},
'executor': executor,
},
show_only=["templates/workers/worker-kedaautoscaler.yaml"],
)
if is_created:
assert jmespath.search("metadata.name",
docs[0]) == "RELEASE-NAME-worker"
else:
assert docs == []
def test_service_account_custom_names(self):
k8s_objects = render_chart(
"TEST-RBAC",
values={
"fullnameOverride": "TEST-RBAC",
"cleanup": {
"enabled": True,
"serviceAccount": {
"name": CUSTOM_CLEANUP_NAME,
},
},
"scheduler": {"serviceAccount": {"name": CUSTOM_SCHEDULER_NAME}},
"webserver": {"serviceAccount": {"name": CUSTOM_WEBSERVER_NAME}},
"workers": {"serviceAccount": {"name": CUSTOM_WORKER_NAME}},
"flower": {"serviceAccount": {"name": CUSTOM_FLOWER_NAME}},
"statsd": {"serviceAccount": {"name": CUSTOM_STATSD_NAME}},
"redis": {"serviceAccount": {"name": CUSTOM_REDIS_NAME}},
"pgbouncer": {
"enabled": True,
"serviceAccount": {
"name": CUSTOM_PGBOUNCER_NAME,
},
},
"createUserJob": {"serviceAccount": {"name": CUSTOM_CREATE_USER_JOBS_NAME}},
"migrateDatabaseJob": {"serviceAccount": {"name": CUSTOM_MIGRATE_DATABASE_JOBS_NAME}},
},
)
list_of_sa_names = [
k8s_object['metadata']['name']
for k8s_object in k8s_objects
if k8s_object['kind'] == "ServiceAccount"
]
self.assertCountEqual(
list_of_sa_names,
CUSTOM_SERVICE_ACCOUNT_NAMES,
)
def test_deployments_no_rbac_no_sa(self):
k8s_objects = render_chart(
"TEST-RBAC",
values={
"fullnameOverride": "TEST-RBAC",
"rbac": {"create": False},
"cleanup": {
"enabled": True,
"serviceAccount": {
"create": False,
},
},
"pgbouncer": {
"enabled": True,
"serviceAccount": {
"create": False,
},
},
"redis": {"serviceAccount": {"create": False}},
"scheduler": {"serviceAccount": {"create": False}},
"webserver": {"serviceAccount": {"create": False}},
"workers": {"serviceAccount": {"create": False}},
"statsd": {"serviceAccount": {"create": False}},
"createUserJob": {"serviceAccount": {"create": False}},
"migrateDatabaseJob": {"serviceAccount": {"create": False}},
"flower": {"serviceAccount": {"create": False}},
},
)
list_of_kind_names_tuples = [
(k8s_object['kind'], k8s_object['metadata']['name']) for k8s_object in k8s_objects
]
self.assertCountEqual(
list_of_kind_names_tuples,
DEPLOYMENT_NO_RBAC_NO_SA_KIND_NAME_TUPLES,
)
def test_should_add_dags_volume(self):
docs = render_chart(
values={"dags": {
"gitSync": {
"enabled": True
}
}},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
# check that there is a volume and git-sync and scheduler container mount it
assert len(
jmespath.search("spec.template.spec.volumes[?name=='dags']",
docs[0])) > 0
assert (len(
jmespath.search(
"(spec.template.spec.containers[?name=='scheduler'].volumeMounts[])[?name=='dags']",
docs[0],
)) > 0)
assert (len(
jmespath.search(
"(spec.template.spec.containers[?name=='git-sync'].volumeMounts[])[?name=='dags']",
docs[0],
)) > 0)
def test_validate_if_ssh_params_are_added(self):
docs = render_chart(
values={
"dags": {
"gitSync": {
"enabled": True,
"containerName": "git-sync-test",
"sshKeySecret": "ssh-secret",
"knownHosts": None,
"branch": "test-branch",
}
}
},
show_only=["templates/pod-template-file.yaml"],
chart_dir=self.temp_chart_dir,
)
assert {
"name": "GIT_SSH_KEY_FILE",
"value": "/etc/git-secret/ssh"
} in jmespath.search("spec.initContainers[0].env", docs[0])
assert {
"name": "GIT_SYNC_SSH",
"value": "true"
} in jmespath.search("spec.initContainers[0].env", docs[0])
assert {
"name": "GIT_KNOWN_HOSTS",
"value": "false"
} in jmespath.search("spec.initContainers[0].env", docs[0])
assert {
"name": "git-sync-ssh-key",
"secret": {
"secretName": "ssh-secret",
"defaultMode": 288
},
} in jmespath.search("spec.volumes", docs[0])
def test_external_redis_broker_url(self, executor):
k8s_objects = render_chart(
RELEASE_NAME_REDIS,
{
"executor": executor,
"networkPolicies": {"enabled": True},
"data": {
"brokerUrl": "redis://*****:*****@redis-host:6379/0",
},
"redis": {"enabled": False},
},
)
k8s_obj_by_key = prepare_k8s_lookup_dict(k8s_objects)
created_redis_objects = SET_POSSIBLE_REDIS_OBJECT_KEYS & set(k8s_obj_by_key.keys())
assert created_redis_objects == {REDIS_OBJECTS["SECRET_BROKER_URL"]}
self.assert_password_and_broker_url_secrets(
k8s_obj_by_key,
expected_password_match=None,
expected_broker_url_match="redis://*****:*****@redis-host:6379/0",
)
self.assert_broker_url_env(k8s_obj_by_key)
def test_git_sync_with_different_airflow_versions(self, airflow_version, exclude_webserver):
"""
If Airflow >= 2.0.0 - git sync related containers, volume mounts & volumes
are not created.
"""
docs = render_chart(
values={
"airflowVersion": airflow_version,
"dags": {
"gitSync": {
"enabled": True,
},
"persistence": {"enabled": False},
},
},
show_only=["templates/webserver/webserver-deployment.yaml"],
)
containers_names = [
container["name"] for container in jmespath.search("spec.template.spec.containers", docs[0])
]
volume_mount_names = [
vm["name"] for vm in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
]
volume_names = [volume["name"] for volume in jmespath.search("spec.template.spec.volumes", docs[0])]
if exclude_webserver:
assert "git-sync" not in containers_names
assert "dags" not in volume_mount_names
assert "dags" not in volume_names
else:
assert "git-sync" in containers_names
assert "dags" in volume_mount_names
assert "dags" in volume_names
def test_should_add_extra_volume_and_extra_volume_mount(self):
docs = render_chart(
values={
"executor": "CeleryExecutor",
"scheduler": {
"extraVolumes": [{"name": "test-volume", "emptyDir": {}}],
"extraVolumeMounts": [{"name": "test-volume", "mountPath": "/opt/test"}],
},
"dags": {
"gitSync": {
"enabled": True,
}
},
},
show_only=["templates/scheduler/scheduler-deployment.yaml"],
)
self.assertIn(
{"name": "test-volume", "emptyDir": {}}, jmespath.search("spec.template.spec.volumes", docs[0])
)
self.assertIn(
{"name": "test-volume", "mountPath": "/opt/test"},
jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0]),
)