[chart/v1-2x-test] Support image digest in worker pod template image (#69776)

chart/templates/_helpers.yaml

@@ -396,7 +396,14 @@ If release name contains chart name it will be used as a full name.
 {{- end }}
 
 {{- define "pod_template_image" -}}
-  {{- printf "%s:%s" (.Values.images.pod_template.repository | default .Values.defaultAirflowRepository) (.Values.images.pod_template.tag | default .Values.defaultAirflowTag) }}
+  {{- $repository := .Values.images.pod_template.repository | default .Values.defaultAirflowRepository -}}
+  {{- $tag := .Values.images.pod_template.tag | default .Values.defaultAirflowTag -}}
+  {{- $digest := .Values.images.pod_template.digest | default .Values.defaultAirflowDigest -}}
+  {{- if $digest }}
+    {{- printf "%s@%s" $repository $digest -}}
+  {{- else }}
+    {{- printf "%s:%s" $repository $tag -}}
+  {{- end }}
 {{- end }}
 
 {{/* This helper is used for airflow containers that do not need the users code */}}

chart/values.schema.json

@@ -920,6 +920,14 @@
                             ],
                             "default": null
                         },
+                        "digest": {
+                            "description": "The pod_template image digest. If set, it will override the tag. If ``config.kubernetes_executor.worker_container_repository`` and ``config.kubernetes_executor.worker_container_tag`` are set, the k8s executor will use those instead.",
+                            "type": [
+                                "string",
+                                "null"
+                            ],
+                            "default": null
+                        },
                         "pullPolicy": {
                             "description": "The pod_template image pull policy.",
                             "type": "string",

chart/values.yaml

@@ -89,12 +89,14 @@ images:
   # timeout (in seconds) for airflow-migrations to complete
   migrationsWaitTimeout: 60
   pod_template:
-    # Note that `images.pod_template.repository` and `images.pod_template.tag` parameters can be overridden
-    # in `config.kubernetes_executor` section. So for these parameters to have effect
-    # `config.kubernetes_executor.worker_container_repository` and
+    # Note that `images.pod_template.repository`, `images.pod_template.tag` and `images.pod_template.digest`
+    # parameters can be overridden in `config.kubernetes_executor` section. So for these parameters
+    # to have effect `config.kubernetes_executor.worker_container_repository` and
     # `config.kubernetes_executor.worker_container_tag` must be not set .
     repository: ~
     tag: ~
+    # Specifying digest takes precedence over tag.
+    digest: ~
     pullPolicy: IfNotPresent
   flower:
     repository: ~

helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py

@@ -357,6 +357,39 @@ def test_should_set_a_custom_image_in_pod_template(self):
         assert jmespath.search("spec.containers[0].imagePullPolicy", docs[0]) == "Always"
         assert jmespath.search("spec.containers[0].name", docs[0]) == "base"
 
+    @pytest.mark.parametrize(
+        ("expected_image", "tag", "digest"),
+        [
+            ("dummy_image:user-tag", "user-tag", None),
+            ("dummy_image@user-digest", None, "user-digest"),
+            ("dummy_image@user-digest", "user-tag", "user-digest"),
+        ],
+    )
+    def test_should_use_correct_pod_template_image(self, expected_image, tag, digest):
+        docs = render_chart(
+            values={
+                "images": {
+                    "pod_template": {"repository": "dummy_image", "tag": tag, "digest": digest},
+                },
+            },
+            show_only=["templates/pod-template-file.yaml"],
+            chart_dir=self.temp_chart_dir,
+        )
+
+        assert jmespath.search("spec.containers[0].image", docs[0]) == expected_image
+
+    def test_should_use_default_airflow_digest_in_pod_template(self):
+        docs = render_chart(
+            values={
+                "defaultAirflowRepository": "test-repo",
+                "defaultAirflowDigest": "user-digest",
+            },
+            show_only=["templates/pod-template-file.yaml"],
+            chart_dir=self.temp_chart_dir,
+        )
+
+        assert jmespath.search("spec.containers[0].image", docs[0]) == "test-repo@user-digest"
+
     def test_mount_airflow_cfg(self):
         docs = render_chart(
             values={},