From 8b0da7701f55626439960614d0f1f7e9866c19fe Mon Sep 17 00:00:00 2001 From: Weichung Shaw Date: Thu, 23 Sep 2021 14:17:40 +0100 Subject: [PATCH 1/3] allow per-component custom pod annotations --- chart/templates/flower/flower-deployment.yaml | 4 +- .../scheduler/scheduler-deployment.yaml | 3 + .../triggerer/triggerer-deployment.yaml | 3 + .../webserver/webserver-deployment.yaml | 3 + .../templates/workers/worker-deployment.yaml | 3 + chart/tests/test_airflow_common.py | 8 +- chart/tests/test_annotations.py | 117 ++++++++++++++++++ chart/values.schema.json | 25 ++++ chart/values.yaml | 11 ++ 9 files changed, 173 insertions(+), 4 deletions(-) diff --git a/chart/templates/flower/flower-deployment.yaml b/chart/templates/flower/flower-deployment.yaml index cb162414174f..a5e145525377 100644 --- a/chart/templates/flower/flower-deployment.yaml +++ b/chart/templates/flower/flower-deployment.yaml @@ -54,9 +54,9 @@ spec: {{- with .Values.labels }} {{ toYaml . | indent 8 }} {{- end }} - {{- if .Values.airflowPodAnnotations }} + {{- if or (.Values.airflowPodAnnotations) (.Values.flower.podAnnotations) }} annotations: - {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} + {{- mustMerge .Values.flower.podAnnotations .Values.airflowPodAnnotations | toYaml | nindent 8 }} {{- end }} spec: nodeSelector: diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml index 0ac86f4aa0b5..a73c6b90e4ce 100644 --- a/chart/templates/scheduler/scheduler-deployment.yaml +++ b/chart/templates/scheduler/scheduler-deployment.yaml @@ -85,6 +85,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.scheduler.podAnnotations }} + {{- toYaml .Values.scheduler.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{ toYaml $nodeSelector | indent 8 }} diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml index fbfbcd279558..e97f45a822a0 100644 --- a/chart/templates/triggerer/triggerer-deployment.yaml +++ b/chart/templates/triggerer/triggerer-deployment.yaml @@ -67,6 +67,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.triggerer.podAnnotations }} + {{- toYaml .Values.triggerer.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{- toYaml $nodeSelector | nindent 8 }} diff --git a/chart/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml index da89222d8410..57e564d421c9 100644 --- a/chart/templates/webserver/webserver-deployment.yaml +++ b/chart/templates/webserver/webserver-deployment.yaml @@ -81,6 +81,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.webserver.podAnnotations }} + {{- toYaml .Values.webserver.podAnnotations | nindent 8 }} + {{- end }} spec: serviceAccountName: {{ include "webserver.serviceAccountName" . }} nodeSelector: diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 5af41dc1e0f4..de6b4c64622a 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -77,6 +77,9 @@ spec: {{- if .Values.airflowPodAnnotations }} {{- toYaml .Values.airflowPodAnnotations | nindent 8 }} {{- end }} + {{- if .Values.workers.podAnnotations }} + {{- toYaml .Values.workers.podAnnotations | nindent 8 }} + {{- end }} spec: nodeSelector: {{ toYaml $nodeSelector | indent 8 }} diff --git a/chart/tests/test_airflow_common.py b/chart/tests/test_airflow_common.py index 9122d35d2d87..66244be5fa5e 100644 --- a/chart/tests/test_airflow_common.py +++ b/chart/tests/test_airflow_common.py @@ -74,16 +74,20 @@ def test_annotations(self): release_name = "TEST-BASIC" k8s_objects = render_chart( name=release_name, - values={"airflowPodAnnotations": {"test-annotation/safe-to-evict": "true"}}, + values={ + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "airflowPodAnnotations": {"test-annotation/safe-to-evict": "true"}, + }, show_only=[ "templates/scheduler/scheduler-deployment.yaml", "templates/workers/worker-deployment.yaml", "templates/webserver/webserver-deployment.yaml", "templates/flower/flower-deployment.yaml", + "templates/triggerer/triggerer-deployment.yaml", ], ) - assert 4 == len(k8s_objects) + assert 5 == len(k8s_objects) for k8s_object in k8s_objects: annotations = k8s_object["spec"]["template"]["metadata"]["annotations"] diff --git a/chart/tests/test_annotations.py b/chart/tests/test_annotations.py index 899d697131a5..0cd3d100b0b4 100644 --- a/chart/tests/test_annotations.py +++ b/chart/tests/test_annotations.py @@ -45,6 +45,14 @@ "redis", ) +COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS = ( + "scheduler", + "webserver", + "worker", + "flower", + "triggerer", +) + class AnnotationsTest(unittest.TestCase): def test_service_account_annotations(self): @@ -136,3 +144,112 @@ def test_service_account_annotations(self): list_of_annotation_values_in_objects, CUSTOM_ANNOTATION_VALUES, ) + + def test_per_component_custom_annotations(self): + release_name = "RELEASE_NAME" + + k8s_objects = render_chart( + name=release_name, + values={ + "scheduler": { + "podAnnotations": { + "example": "scheduler", + }, + }, + "webserver": { + "podAnnotations": { + "example": "webserver", + }, + }, + "workers": { + "podAnnotations": { + "example": "worker", + }, + }, + "flower": { + "podAnnotations": { + "example": "flower", + }, + }, + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "podAnnotations": { + "example": "triggerer", + }, + }, + }, + show_only=[ + "templates/scheduler/scheduler-deployment.yaml", + "templates/workers/worker-deployment.yaml", + "templates/webserver/webserver-deployment.yaml", + "templates/flower/flower-deployment.yaml", + "templates/triggerer/triggerer-deployment.yaml", + ], + ) + + # The test relies on the convention that the Deployment name + # is always `{ Release.Name }-`. + obj_by_component_name = { + obj["metadata"]["name"].replace(release_name + "-", ""): obj for obj in k8s_objects + } + + self.assertCountEqual(obj_by_component_name, COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS) + + for component_name, obj in obj_by_component_name.items(): + self.assertIn("example", obj["spec"]["template"]["metadata"]["annotations"]) + self.assertEqual(component_name, obj["spec"]["template"]["metadata"]["annotations"]["example"]) + + def test_per_component_custom_annotations_precedence(self): + release_name = "RELEASE_NAME" + + k8s_objects = render_chart( + name=release_name, + values={ + "airflowPodAnnotations": {"example": "GLOBAL"}, + "scheduler": { + "podAnnotations": { + "example": "scheduler", + }, + }, + "webserver": { + "podAnnotations": { + "example": "webserver", + }, + }, + "workers": { + "podAnnotations": { + "example": "worker", + }, + }, + "flower": { + "podAnnotations": { + "example": "flower", + }, + }, + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "podAnnotations": { + "example": "triggerer", + }, + }, + }, + show_only=[ + "templates/scheduler/scheduler-deployment.yaml", + "templates/workers/worker-deployment.yaml", + "templates/webserver/webserver-deployment.yaml", + "templates/flower/flower-deployment.yaml", + "templates/triggerer/triggerer-deployment.yaml", + ], + ) + + # The test relies on the convention that the Deployment name + # is always `{ Release.Name }-`. + obj_by_component_name = { + obj["metadata"]["name"].replace(release_name + "-", ""): obj for obj in k8s_objects + } + + self.assertCountEqual(obj_by_component_name, COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS) + + for component_name, obj in obj_by_component_name.items(): + self.assertIn("example", obj["spec"]["template"]["metadata"]["annotations"]) + self.assertEqual(component_name, obj["spec"]["template"]["metadata"]["annotations"]["example"]) diff --git a/chart/values.schema.json b/chart/values.schema.json index 4c01287ca302..012fc2d4d047 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -1141,6 +1141,11 @@ } ] }, + "podAnnotations": { + "description": "Annotations to add to the worker pods.", + "type": "object", + "default": {} + }, "logGroomerSidecar": { "description": "Configuration for worker log groomer sidecar", "type": "object", @@ -1384,6 +1389,11 @@ "type": "array", "default": [] }, + "podAnnotations": { + "description": "Annotations to add to the scheduler pods.", + "type": "object", + "default": {} + }, "logGroomerSidecar": { "description": "Configuration for the schedulers log groomer sidecar.", "type": "object", @@ -1609,6 +1619,11 @@ "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes.", "type": "integer", "default": 60 + }, + "podAnnotations": { + "description": "Annotations to add to the triggerer pods.", + "type": "object", + "default": {} } } }, @@ -2067,6 +2082,11 @@ "description": "Specify Tolerations for webserver pods.", "type": "array", "default": [] + }, + "podAnnotations": { + "description": "Annotations to add to the webserver pods.", + "type": "object", + "default": {} } } }, @@ -2287,6 +2307,11 @@ "description": "Specify Tolerations for Flower pods.", "type": "array", "default": [] + }, + "podAnnotations": { + "description": "Annotations to add to the Flower pods.", + "type": "object", + "default": {} } } }, diff --git a/chart/values.yaml b/chart/values.yaml index 0bc344c8098b..b535edc48330 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -485,6 +485,8 @@ workers: # hostnames: # - "test.hostname.two" + podAnnotations: {} + logGroomerSidecar: # Command to use when running the Airflow worker log groomer sidecar (templated). command: ~ @@ -578,6 +580,8 @@ scheduler: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + logGroomerSidecar: # Whether to deploy the Airflow scheduler log groomer sidecar. enabled: true @@ -594,6 +598,7 @@ scheduler: # requests: # cpu: 100m # memory: 128Mi + # Airflow create user job settings createUserJob: # Annotations on the create user job pod @@ -762,6 +767,8 @@ webserver: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + # Airflow Triggerer Config triggerer: # Number of airflow triggerers in the deployment @@ -834,6 +841,8 @@ triggerer: topologyKey: "kubernetes.io/hostname" tolerations: [] + podAnnotations: {} + # Flower settings flower: # Enable flower. @@ -915,6 +924,8 @@ flower: affinity: {} tolerations: [] + podAnnotations: {} + # Statsd settings statsd: enabled: true From 2eb6b9ebbfa48f4de856ef81ec759880fc1cc485 Mon Sep 17 00:00:00 2001 From: Weichung Shaw Date: Thu, 23 Sep 2021 19:55:49 +0100 Subject: [PATCH 2/3] refactor annotations test for readability --- chart/tests/test_annotations.py | 46 +++++++++++++++------------------ 1 file changed, 21 insertions(+), 25 deletions(-) diff --git a/chart/tests/test_annotations.py b/chart/tests/test_annotations.py index 0cd3d100b0b4..095bd44b5933 100644 --- a/chart/tests/test_annotations.py +++ b/chart/tests/test_annotations.py @@ -19,20 +19,7 @@ from tests.helm_template_generator import render_chart -# Values for each service mapped to the 'example' -# key annotation -CUSTOM_ANNOTATION_VALUES = ( - CUSTOM_SCHEDULER_ANNOTATION, - CUSTOM_WEBSERVER_ANNOTATION, - CUSTOM_WORKER_ANNOTATION, - CUSTOM_CLEANUP_ANNOTATION, - CUSTOM_FLOWER_ANNOTATION, - CUSTOM_PGBOUNCER_ANNOTATION, - CUSTOM_STATSD_ANNOTATION, - CUSTOM_CREATE_USER_JOB_ANNOTATION, - CUSTOM_MIGRATE_DATABASE_JOB_ANNOTATION, - CUSTOM_REDIS_ANNOTATION, -) = ( +COMPONENTS_SUPPORTING_CUSTOM_SERVICEACCOUNT_ANNOTATIONS = ( "scheduler", "webserver", "worker", @@ -43,6 +30,7 @@ "createuser", "migratedb", "redis", + "triggerer", ) COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS = ( @@ -62,49 +50,49 @@ def test_service_account_annotations(self): "enabled": True, "serviceAccount": { "annotations": { - "example": CUSTOM_CLEANUP_ANNOTATION, + "example": "cleanup", }, }, }, "scheduler": { "serviceAccount": { "annotations": { - "example": CUSTOM_SCHEDULER_ANNOTATION, + "example": "scheduler", }, }, }, "webserver": { "serviceAccount": { "annotations": { - "example": CUSTOM_WEBSERVER_ANNOTATION, + "example": "webserver", }, }, }, "workers": { "serviceAccount": { "annotations": { - "example": CUSTOM_WORKER_ANNOTATION, + "example": "worker", }, }, }, "flower": { "serviceAccount": { "annotations": { - "example": CUSTOM_FLOWER_ANNOTATION, + "example": "flower", }, }, }, "statsd": { "serviceAccount": { "annotations": { - "example": CUSTOM_STATSD_ANNOTATION, + "example": "statsd", }, }, }, "redis": { "serviceAccount": { "annotations": { - "example": CUSTOM_REDIS_ANNOTATION, + "example": "redis", }, }, }, @@ -112,25 +100,33 @@ def test_service_account_annotations(self): "enabled": True, "serviceAccount": { "annotations": { - "example": CUSTOM_PGBOUNCER_ANNOTATION, + "example": "pgbouncer", }, }, }, "createUserJob": { "serviceAccount": { "annotations": { - "example": CUSTOM_CREATE_USER_JOB_ANNOTATION, + "example": "createuser", }, }, }, "migrateDatabaseJob": { "serviceAccount": { "annotations": { - "example": CUSTOM_MIGRATE_DATABASE_JOB_ANNOTATION, + "example": "migratedb", }, }, }, "executor": "CeleryExecutor", # create worker deployment + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "serviceAccount": { + "annotations": { + "example": "triggerer", + }, + }, + }, }, ) @@ -142,7 +138,7 @@ def test_service_account_annotations(self): self.assertCountEqual( list_of_annotation_values_in_objects, - CUSTOM_ANNOTATION_VALUES, + COMPONENTS_SUPPORTING_CUSTOM_SERVICEACCOUNT_ANNOTATIONS, ) def test_per_component_custom_annotations(self): From 6b9861ec43f711888d15c36a0980c7b810c7530f Mon Sep 17 00:00:00 2001 From: Weichung Shaw Date: Fri, 24 Sep 2021 17:45:57 +0100 Subject: [PATCH 3/3] more test_annotations refactor --- chart/tests/test_annotations.py | 382 +++++++++++++++++++------------- 1 file changed, 227 insertions(+), 155 deletions(-) diff --git a/chart/tests/test_annotations.py b/chart/tests/test_annotations.py index 095bd44b5933..48dd28f688f1 100644 --- a/chart/tests/test_annotations.py +++ b/chart/tests/test_annotations.py @@ -15,158 +15,260 @@ # specific language governing permissions and limitations # under the License. -import unittest +import pytest from tests.helm_template_generator import render_chart -COMPONENTS_SUPPORTING_CUSTOM_SERVICEACCOUNT_ANNOTATIONS = ( - "scheduler", - "webserver", - "worker", - "cleanup", - "flower", - "pgbouncer", - "statsd", - "createuser", - "migratedb", - "redis", - "triggerer", -) - -COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS = ( - "scheduler", - "webserver", - "worker", - "flower", - "triggerer", -) - -class AnnotationsTest(unittest.TestCase): - def test_service_account_annotations(self): - k8s_objects = render_chart( - values={ - "cleanup": { - "enabled": True, - "serviceAccount": { - "annotations": { - "example": "cleanup", +class TestServiceAccountAnnotations: + @pytest.mark.parametrize( + "values,show_only,expected_annotations", + [ + ( + { + "cleanup": { + "enabled": True, + "serviceAccount": { + "annotations": { + "example": "cleanup", + }, }, }, }, - "scheduler": { - "serviceAccount": { - "annotations": { - "example": "scheduler", + "templates/cleanup/cleanup-serviceaccount.yaml", + { + "example": "cleanup", + }, + ), + ( + { + "scheduler": { + "serviceAccount": { + "annotations": { + "example": "scheduler", + }, }, }, }, - "webserver": { - "serviceAccount": { - "annotations": { - "example": "webserver", + "templates/scheduler/scheduler-serviceaccount.yaml", + { + "example": "scheduler", + }, + ), + ( + { + "webserver": { + "serviceAccount": { + "annotations": { + "example": "webserver", + }, }, }, }, - "workers": { - "serviceAccount": { - "annotations": { - "example": "worker", + "templates/webserver/webserver-serviceaccount.yaml", + { + "example": "webserver", + }, + ), + ( + { + "workers": { + "serviceAccount": { + "annotations": { + "example": "worker", + }, }, }, }, - "flower": { - "serviceAccount": { - "annotations": { - "example": "flower", + "templates/workers/worker-serviceaccount.yaml", + { + "example": "worker", + }, + ), + ( + { + "flower": { + "serviceAccount": { + "annotations": { + "example": "flower", + }, }, }, }, - "statsd": { - "serviceAccount": { - "annotations": { - "example": "statsd", + "templates/flower/flower-serviceaccount.yaml", + { + "example": "flower", + }, + ), + ( + { + "statsd": { + "serviceAccount": { + "annotations": { + "example": "statsd", + }, }, }, }, - "redis": { - "serviceAccount": { - "annotations": { - "example": "redis", + "templates/statsd/statsd-serviceaccount.yaml", + { + "example": "statsd", + }, + ), + ( + { + "redis": { + "serviceAccount": { + "annotations": { + "example": "redis", + }, }, }, }, - "pgbouncer": { - "enabled": True, - "serviceAccount": { - "annotations": { - "example": "pgbouncer", + "templates/redis/redis-serviceaccount.yaml", + { + "example": "redis", + }, + ), + ( + { + "pgbouncer": { + "enabled": True, + "serviceAccount": { + "annotations": { + "example": "pgbouncer", + }, }, }, }, - "createUserJob": { - "serviceAccount": { - "annotations": { - "example": "createuser", + "templates/pgbouncer/pgbouncer-serviceaccount.yaml", + { + "example": "pgbouncer", + }, + ), + ( + { + "createUserJob": { + "serviceAccount": { + "annotations": { + "example": "createuser", + }, }, }, }, - "migrateDatabaseJob": { - "serviceAccount": { - "annotations": { - "example": "migratedb", + "templates/jobs/create-user-job-serviceaccount.yaml", + { + "example": "createuser", + }, + ), + ( + { + "migrateDatabaseJob": { + "serviceAccount": { + "annotations": { + "example": "migratedb", + }, }, }, }, - "executor": "CeleryExecutor", # create worker deployment - "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. - "triggerer": { - "serviceAccount": { - "annotations": { - "example": "triggerer", + "templates/jobs/migrate-database-job-serviceaccount.yaml", + { + "example": "migratedb", + }, + ), + ( + { + "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. + "triggerer": { + "serviceAccount": { + "annotations": { + "example": "triggerer", + }, }, }, }, - }, + "templates/triggerer/triggerer-serviceaccount.yaml", + { + "example": "triggerer", + }, + ), + ], + ) + def test_annotations_are_added(self, values, show_only, expected_annotations): + k8s_objects = render_chart( + values=values, + show_only=[show_only], ) - list_of_annotation_values_in_objects = [ - k8s_object['metadata']['annotations']['example'] - for k8s_object in k8s_objects - if k8s_object['kind'] == "ServiceAccount" - ] + # This test relies on the convention that the helm chart puts a single + # ServiceAccount in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] - self.assertCountEqual( - list_of_annotation_values_in_objects, - COMPONENTS_SUPPORTING_CUSTOM_SERVICEACCOUNT_ANNOTATIONS, - ) + for k, v in expected_annotations.items(): + assert k in obj["metadata"]["annotations"] + assert v == obj["metadata"]["annotations"][k] - def test_per_component_custom_annotations(self): - release_name = "RELEASE_NAME" - k8s_objects = render_chart( - name=release_name, - values={ +@pytest.mark.parametrize( + "values,show_only,expected_annotations", + [ + ( + { "scheduler": { "podAnnotations": { "example": "scheduler", }, }, + }, + "templates/scheduler/scheduler-deployment.yaml", + { + "example": "scheduler", + }, + ), + ( + { "webserver": { "podAnnotations": { "example": "webserver", }, }, + }, + "templates/webserver/webserver-deployment.yaml", + { + "example": "webserver", + }, + ), + ( + { "workers": { "podAnnotations": { "example": "worker", }, }, + }, + "templates/workers/worker-deployment.yaml", + { + "example": "worker", + }, + ), + ( + { "flower": { "podAnnotations": { "example": "flower", }, }, + }, + "templates/flower/flower-deployment.yaml", + { + "example": "flower", + }, + ), + ( + { "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. "triggerer": { "podAnnotations": { @@ -174,78 +276,48 @@ def test_per_component_custom_annotations(self): }, }, }, - show_only=[ - "templates/scheduler/scheduler-deployment.yaml", - "templates/workers/worker-deployment.yaml", - "templates/webserver/webserver-deployment.yaml", - "templates/flower/flower-deployment.yaml", - "templates/triggerer/triggerer-deployment.yaml", - ], + "templates/triggerer/triggerer-deployment.yaml", + { + "example": "triggerer", + }, + ), + ], +) +class TestPerComponentPodAnnotations: + def test_annotations_are_added(self, values, show_only, expected_annotations): + k8s_objects = render_chart( + values=values, + show_only=[show_only], ) - # The test relies on the convention that the Deployment name - # is always `{ Release.Name }-`. - obj_by_component_name = { - obj["metadata"]["name"].replace(release_name + "-", ""): obj for obj in k8s_objects - } + # This test relies on the convention that the helm chart puts a single + # Deployment in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] - self.assertCountEqual(obj_by_component_name, COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS) + for k, v in expected_annotations.items(): + assert k in obj["spec"]["template"]["metadata"]["annotations"] + assert v == obj["spec"]["template"]["metadata"]["annotations"][k] - for component_name, obj in obj_by_component_name.items(): - self.assertIn("example", obj["spec"]["template"]["metadata"]["annotations"]) - self.assertEqual(component_name, obj["spec"]["template"]["metadata"]["annotations"]["example"]) + def test_precedence(self, values, show_only, expected_annotations): + values_global_annotations = {"airflowPodAnnotations": {k: "GLOBAL" for k in expected_annotations}} - def test_per_component_custom_annotations_precedence(self): - release_name = "RELEASE_NAME" + values_merged = {**values, **values_global_annotations} k8s_objects = render_chart( - name=release_name, - values={ - "airflowPodAnnotations": {"example": "GLOBAL"}, - "scheduler": { - "podAnnotations": { - "example": "scheduler", - }, - }, - "webserver": { - "podAnnotations": { - "example": "webserver", - }, - }, - "workers": { - "podAnnotations": { - "example": "worker", - }, - }, - "flower": { - "podAnnotations": { - "example": "flower", - }, - }, - "airflowVersion": "2.2.0", # Needed for triggerer to be enabled. - "triggerer": { - "podAnnotations": { - "example": "triggerer", - }, - }, - }, - show_only=[ - "templates/scheduler/scheduler-deployment.yaml", - "templates/workers/worker-deployment.yaml", - "templates/webserver/webserver-deployment.yaml", - "templates/flower/flower-deployment.yaml", - "templates/triggerer/triggerer-deployment.yaml", - ], + values=values_merged, + show_only=[show_only], ) - # The test relies on the convention that the Deployment name - # is always `{ Release.Name }-`. - obj_by_component_name = { - obj["metadata"]["name"].replace(release_name + "-", ""): obj for obj in k8s_objects - } - - self.assertCountEqual(obj_by_component_name, COMPONENTS_SUPPORTING_CUSTOM_POD_ANNOTATIONS) + # This test relies on the convention that the helm chart puts a single + # Deployment in its own .yaml file, so by specifying `show_only`, + # we should only get a single k8s_object here - the target object that + # we hope to test on. + assert len(k8s_objects) == 1 + obj = k8s_objects[0] - for component_name, obj in obj_by_component_name.items(): - self.assertIn("example", obj["spec"]["template"]["metadata"]["annotations"]) - self.assertEqual(component_name, obj["spec"]["template"]["metadata"]["annotations"]["example"]) + for k, v in expected_annotations.items(): + assert k in obj["spec"]["template"]["metadata"]["annotations"] + assert v == obj["spec"]["template"]["metadata"]["annotations"][k]