Miretpl commented on code in PR #63514: URL: https://github.com/apache/airflow/pull/63514#discussion_r2969313602
########## helm-tests/tests/helm_tests/airflow_core/test_worker_image.py: ########## @@ -0,0 +1,693 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import jmespath +import pytest +import yaml +from chart_utils.helm_template_generator import CHART_DIR, render_chart + + +def _get_default_airflow_image() -> tuple[str, str]: + """Read the default airflow image and tag from chart/values.yaml.""" + values = yaml.safe_load((CHART_DIR / "values.yaml").read_text()) + repo = values["defaultAirflowRepository"] + tag = values["defaultAirflowTag"] + return f"{repo}:{tag}", tag + + +DEFAULT_AIRFLOW_IMAGE, DEFAULT_AIRFLOW_TAG = _get_default_airflow_image() + + +def _all_airflow_images(doc): + """Extract all airflow images from a rendered worker deployment/statefulset.""" + init_images = jmespath.search("spec.template.spec.initContainers[*].image", doc) or [] + container_images = jmespath.search("spec.template.spec.containers[*].image", doc) or [] + # Filter out git-sync sidecar images (they don't use airflow_worker_image) + all_images = init_images + container_images + return [img for img in all_images if "git-sync" not in img] + + +def _all_airflow_pull_policies(doc): + """Extract all imagePullPolicy values from airflow containers.""" + init_policies = jmespath.search("spec.template.spec.initContainers[*].imagePullPolicy", doc) or [] + container_policies = jmespath.search("spec.template.spec.containers[*].imagePullPolicy", doc) or [] + return init_policies + container_policies + + +class TestWorkerImageDefault: + """Tests that workers use the default airflow image when no override is set.""" + + def test_default_image_used_when_no_override(self): + """When no worker image override is set, all worker containers use the global default image.""" + docs = render_chart( + values={"executor": "CeleryExecutor"}, + show_only=["templates/workers/worker-deployment.yaml"], + ) + assert len(docs) == 1 + images = _all_airflow_images(docs[0]) + assert len(images) > 0 + for image in images: + assert image == DEFAULT_AIRFLOW_IMAGE + + def test_default_image_with_global_override(self): + """When images.airflow is set, workers use it (no worker-specific override).""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "images": { + "airflow": { + "repository": "my-registry/my-airflow", + "tag": "custom-v1", + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == "my-registry/my-airflow:custom-v1" + + +class TestWorkerImageCeleryOverride: Review Comment: All of the tests regarding `celery` overwrite `workers` are in the general worker test file. Could we move those tests there? ########## helm-tests/tests/helm_tests/airflow_core/test_worker_image.py: ########## @@ -0,0 +1,693 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import jmespath +import pytest +import yaml +from chart_utils.helm_template_generator import CHART_DIR, render_chart + + +def _get_default_airflow_image() -> tuple[str, str]: + """Read the default airflow image and tag from chart/values.yaml.""" + values = yaml.safe_load((CHART_DIR / "values.yaml").read_text()) + repo = values["defaultAirflowRepository"] + tag = values["defaultAirflowTag"] + return f"{repo}:{tag}", tag + + +DEFAULT_AIRFLOW_IMAGE, DEFAULT_AIRFLOW_TAG = _get_default_airflow_image() + + +def _all_airflow_images(doc): + """Extract all airflow images from a rendered worker deployment/statefulset.""" + init_images = jmespath.search("spec.template.spec.initContainers[*].image", doc) or [] + container_images = jmespath.search("spec.template.spec.containers[*].image", doc) or [] + # Filter out git-sync sidecar images (they don't use airflow_worker_image) + all_images = init_images + container_images + return [img for img in all_images if "git-sync" not in img] + + +def _all_airflow_pull_policies(doc): + """Extract all imagePullPolicy values from airflow containers.""" + init_policies = jmespath.search("spec.template.spec.initContainers[*].imagePullPolicy", doc) or [] + container_policies = jmespath.search("spec.template.spec.containers[*].imagePullPolicy", doc) or [] + return init_policies + container_policies + + +class TestWorkerImageDefault: + """Tests that workers use the default airflow image when no override is set.""" + + def test_default_image_used_when_no_override(self): + """When no worker image override is set, all worker containers use the global default image.""" + docs = render_chart( + values={"executor": "CeleryExecutor"}, + show_only=["templates/workers/worker-deployment.yaml"], + ) + assert len(docs) == 1 + images = _all_airflow_images(docs[0]) + assert len(images) > 0 + for image in images: + assert image == DEFAULT_AIRFLOW_IMAGE + + def test_default_image_with_global_override(self): + """When images.airflow is set, workers use it (no worker-specific override).""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "images": { + "airflow": { + "repository": "my-registry/my-airflow", + "tag": "custom-v1", + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == "my-registry/my-airflow:custom-v1" + + +class TestWorkerImageCeleryOverride: + """Tests that workers.celery.image overrides the global airflow image.""" + + def test_celery_image_overrides_default(self): + """Setting workers.celery.image.repository/tag overrides the default for all worker containers.""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "workers": { + "celery": { + "image": { + "repository": "celery-custom/airflow", + "tag": "celery-v1", + }, + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + assert len(images) > 0 + for image in images: + assert image == "celery-custom/airflow:celery-v1" + + def test_celery_image_overrides_global_image(self): + """workers.celery.image takes precedence over images.airflow.""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "images": { + "airflow": { + "repository": "global/airflow", + "tag": "global-v1", + }, + }, + "workers": { + "celery": { + "image": { + "repository": "celery-custom/airflow", + "tag": "celery-v2", + }, + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == "celery-custom/airflow:celery-v2" + + def test_celery_image_partial_override_repository_only(self): + """Setting only repository at celery level falls back to default tag.""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "workers": { + "celery": { + "image": { + "repository": "celery-custom/airflow", + }, + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == f"celery-custom/airflow:{DEFAULT_AIRFLOW_TAG}" + + def test_celery_image_partial_override_tag_only(self): + """Setting only tag at celery level falls back to default repository.""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "workers": { + "celery": { + "image": { + "tag": "custom-tag", + }, + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == "apache/airflow:custom-tag" + + def test_celery_image_digest_takes_precedence_over_tag(self): + """When digest is set, it takes precedence over tag.""" + docs = render_chart( + values={ + "executor": "CeleryExecutor", + "workers": { + "celery": { + "image": { + "repository": "celery-custom/airflow", + "tag": "should-be-ignored", + "digest": "sha256:abcdef1234567890", + }, + }, + }, + }, + show_only=["templates/workers/worker-deployment.yaml"], + ) + images = _all_airflow_images(docs[0]) + for image in images: + assert image == "celery-custom/airflow@sha256:abcdef1234567890" + + +class TestWorkerImagePerSetOverride: Review Comment: All of the tests for `sets` overwrite of `celery`, and the `workers` section is in the workers sets dedicated test file. Could we move these tests there? ########## helm-tests/tests/helm_tests/airflow_core/test_worker_image.py: ########## @@ -0,0 +1,693 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import jmespath +import pytest +import yaml +from chart_utils.helm_template_generator import CHART_DIR, render_chart + + +def _get_default_airflow_image() -> tuple[str, str]: + """Read the default airflow image and tag from chart/values.yaml.""" + values = yaml.safe_load((CHART_DIR / "values.yaml").read_text()) + repo = values["defaultAirflowRepository"] + tag = values["defaultAirflowTag"] + return f"{repo}:{tag}", tag + + +DEFAULT_AIRFLOW_IMAGE, DEFAULT_AIRFLOW_TAG = _get_default_airflow_image() + + +def _all_airflow_images(doc): + """Extract all airflow images from a rendered worker deployment/statefulset.""" + init_images = jmespath.search("spec.template.spec.initContainers[*].image", doc) or [] + container_images = jmespath.search("spec.template.spec.containers[*].image", doc) or [] + # Filter out git-sync sidecar images (they don't use airflow_worker_image) + all_images = init_images + container_images + return [img for img in all_images if "git-sync" not in img] + + +def _all_airflow_pull_policies(doc): + """Extract all imagePullPolicy values from airflow containers.""" + init_policies = jmespath.search("spec.template.spec.initContainers[*].imagePullPolicy", doc) or [] + container_policies = jmespath.search("spec.template.spec.containers[*].imagePullPolicy", doc) or [] + return init_policies + container_policies Review Comment: This added test complexity, which can lead to not-so-straightforward debugging in case of failure. Maybe we could simplify the test cases a little? ########## chart/templates/_helpers.yaml: ########## @@ -385,6 +385,44 @@ If release name contains chart name it will be used as a full name. {{- end }} {{- end }} +{{/* + Worker image helper. Users configure image at Values.workers.celery.image (or per-set via + Values.workers.celery.sets[].image). By the time this template is invoked, workersMergeValues + in worker-deployment.yaml has already merged celery/set values into Values.workers, so the + effective image is at Values.workers.image (post-merge). +*/}} +{{- define "airflow_worker_image" -}} Review Comment: Most of this function copies logic from `airflow_image`. Maybe there is some nice way to not duplicate code here? ########## chart/values.yaml: ########## @@ -1148,13 +1148,26 @@ workers: # Queue name for the default workers queue: "default" + # Override the Airflow image for Celery workers. + # If not set, falls back to the global `images.airflow` image. + # Can also be overridden per worker set in `workers.celery.sets[].image`. Review Comment: ```suggestion ``` Not sure if needed. Every section can be overwritten by sets and direct message for one field could lead to `this has this information and other field do not, so I can overwrite this one but this one not?` understanding. ########## helm-tests/tests/helm_tests/airflow_core/test_worker_image.py: ########## @@ -0,0 +1,693 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import jmespath +import pytest +import yaml +from chart_utils.helm_template_generator import CHART_DIR, render_chart + + +def _get_default_airflow_image() -> tuple[str, str]: + """Read the default airflow image and tag from chart/values.yaml.""" + values = yaml.safe_load((CHART_DIR / "values.yaml").read_text()) + repo = values["defaultAirflowRepository"] + tag = values["defaultAirflowTag"] + return f"{repo}:{tag}", tag + + +DEFAULT_AIRFLOW_IMAGE, DEFAULT_AIRFLOW_TAG = _get_default_airflow_image() + + +def _all_airflow_images(doc): + """Extract all airflow images from a rendered worker deployment/statefulset.""" + init_images = jmespath.search("spec.template.spec.initContainers[*].image", doc) or [] + container_images = jmespath.search("spec.template.spec.containers[*].image", doc) or [] + # Filter out git-sync sidecar images (they don't use airflow_worker_image) + all_images = init_images + container_images + return [img for img in all_images if "git-sync" not in img] + + +def _all_airflow_pull_policies(doc): + """Extract all imagePullPolicy values from airflow containers.""" + init_policies = jmespath.search("spec.template.spec.initContainers[*].imagePullPolicy", doc) or [] + container_policies = jmespath.search("spec.template.spec.containers[*].imagePullPolicy", doc) or [] + return init_policies + container_policies + + +class TestWorkerImageDefault: Review Comment: Not sure where the `images` section tests are stored, but I think we should move them there to have all `images` test in one place. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
