Chart: Add annotations to cleanup pods (#21484)

apache · Feb 11, 2022 · 8c1512b · 8c1512b
1 parent c25534b
commit 8c1512b
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 6 deletions.
diff --git a/chart/templates/cleanup/cleanup-cronjob.yaml b/chart/templates/cleanup/cleanup-cronjob.yaml
@@ -58,6 +58,12 @@ spec:
 {{- end }}
           annotations:
             sidecar.istio.io/inject: "false"
+            {{- if .Values.airflowPodAnnotations }}
+            {{- toYaml .Values.airflowPodAnnotations | nindent 12 }}
+            {{- end }}
+            {{- if .Values.cleanup.podAnnotations }}
+            {{- toYaml .Values.cleanup.podAnnotations | nindent 12 }}
+            {{- end }}
         spec:
           restartPolicy: Never
           nodeSelector:

diff --git a/chart/tests/test_airflow_common.py b/chart/tests/test_airflow_common.py
@@ -76,20 +76,26 @@ def test_annotations(self):
             name=release_name,
             values={
                 "airflowPodAnnotations": {"test-annotation/safe-to-evict": "true"},
+                "cleanup": {"enabled": 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",
+                "templates/cleanup/cleanup-cronjob.yaml",
             ],
         )
 
-        assert 5 == len(k8s_objects)
+        assert 6 == len(k8s_objects)
 
         for k8s_object in k8s_objects:
-            annotations = k8s_object["spec"]["template"]["metadata"]["annotations"]
+            if k8s_object['kind'] == 'CronJob':
+                annotations = k8s_object["spec"]["jobTemplate"]["spec"]["template"]["metadata"]["annotations"]
+            else:
+                annotations = k8s_object["spec"]["template"]["metadata"]["annotations"]
+
             assert "test-annotation/safe-to-evict" in annotations
             assert "true" in annotations["test-annotation/safe-to-evict"]
 

diff --git a/chart/tests/test_annotations.py b/chart/tests/test_annotations.py
@@ -20,6 +20,20 @@
 from tests.helm_template_generator import render_chart
 
 
+def deployment_annotations(obj):
+    return obj["spec"]["template"]["metadata"]["annotations"]
+
+
+def cronjob_annotations(obj):
+    return obj["spec"]["jobTemplate"]["spec"]["template"]["metadata"]["annotations"]
+
+
+def get_object_annotations(obj):
+    if obj["kind"] == "CronJob":
+        return cronjob_annotations(obj)
+    return deployment_annotations(obj)
+
+
 class TestServiceAccountAnnotations:
     @pytest.mark.parametrize(
         "values,show_only,expected_annotations",
@@ -279,6 +293,20 @@ def test_annotations_are_added(self, values, show_only, expected_annotations):
                 "example": "triggerer",
             },
         ),
+        (
+            {
+                "cleanup": {
+                    "enabled": True,
+                    "podAnnotations": {
+                        "example": "cleanup",
+                    },
+                }
+            },
+            "templates/cleanup/cleanup-cronjob.yaml",
+            {
+                "example": "cleanup",
+            },
+        ),
     ],
 )
 class TestPerComponentPodAnnotations:
@@ -294,10 +322,11 @@ def test_annotations_are_added(self, values, show_only, expected_annotations):
         # we hope to test on.
         assert len(k8s_objects) == 1
         obj = k8s_objects[0]
+        annotations = get_object_annotations(obj)
 
         for k, v in expected_annotations.items():
-            assert k in obj["spec"]["template"]["metadata"]["annotations"]
-            assert v == obj["spec"]["template"]["metadata"]["annotations"][k]
+            assert k in annotations
+            assert v == annotations[k]
 
     def test_precedence(self, values, show_only, expected_annotations):
         values_global_annotations = {"airflowPodAnnotations": {k: "GLOBAL" for k in expected_annotations}}
@@ -315,7 +344,8 @@ def test_precedence(self, values, show_only, expected_annotations):
         # we hope to test on.
         assert len(k8s_objects) == 1
         obj = k8s_objects[0]
+        annotations = get_object_annotations(obj)
 
         for k, v in expected_annotations.items():
-            assert k in obj["spec"]["template"]["metadata"]["annotations"]
-            assert v == obj["spec"]["template"]["metadata"]["annotations"][k]
+            assert k in annotations
+            assert v == annotations[k]
diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -3846,6 +3846,14 @@
                         "$ref": "#/definitions/io.k8s.api.core.v1.Toleration"
                     }
                 },
+                "podAnnotations": {
+                    "description": "Annotations to add to cleanup pods.",
+                    "type": "object",
+                    "default": {},
+                    "additionalProperties": {
+                        "type": "string"
+                    }
+                },
                 "resources": {
                     "description": "Resources for or cleanup pods",
                     "type": "object",

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -1330,6 +1330,8 @@ cleanup:
   affinity: {}
   tolerations: []
 
+  podAnnotations: {}
+
   resources: {}
   #  limits:
   #   cpu: 100m