From 8b0da7701f55626439960614d0f1f7e9866c19fe Mon Sep 17 00:00:00 2001
From: Weichung Shaw <weichung.shaw@gmail.com>
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 cb162414174f5..a5e145525377d 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 0ac86f4aa0b57..a73c6b90e4ced 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 fbfbcd2795583..e97f45a822a09 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 da89222d84105..57e564d421c91 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 5af41dc1e0f40..de6b4c64622a3 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 9122d35d2d87c..66244be5fa5e3 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 899d697131a53..0cd3d100b0b43 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 }-<component_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 }-<component_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 4c01287ca3028..012fc2d4d0478 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 0bc344c8098ba..b535edc48330d 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 <weichung.shaw@gmail.com>
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 0cd3d100b0b43..095bd44b59331 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 <weichung.shaw@gmail.com>
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 095bd44b59331..48dd28f688f1e 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 }-<component_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 }-<component_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]
