Parallelize Helm tests with multiple job runners (#30672)

Helm Unit tests are using template rendering and the rendering
uses a lot of CPU for `helm template command`. We have a lot of
those rendering tests (>800) so even running the tests in parallel
on a multi-cpu machine does not lead to a decreased elapsed time
to execute the tests.

However, each of the tests is run entirely independently and we
should be able to achieve much faster elapsed time if we run
a subset of tetsts on separate, multi-CPU machine. This will not
lower the job build time, however it might speed up elapsed time
and thus get a faster feedback.

This PR achieves that.

(cherry picked from commit 731ef3d)

.github/workflows/ci.yml

@@ -167,6 +167,7 @@ jobs:
  default-constraints-branch: ${{ steps.selective-checks.outputs.default-constraints-branch }}
  docs-filter: ${{ steps.selective-checks.outputs.docs-filter }}
  skip-pre-commits: ${{ steps.selective-checks.outputs.skip-pre-commits }}
+ helm-test-packages: ${{ steps.selective-checks.outputs.helm-test-packages }}
  debug-resources: ${{ steps.selective-checks.outputs.debug-resources }}
  suspended-providers-folders: ${{ steps.selective-checks.outputs.suspended-providers-folders }}
  source-head-repo: ${{ steps.source-run-info.outputs.source-head-repo }}
@@ -809,9 +810,13 @@ jobs:
 
  tests-helm:
  timeout-minutes: 80
- name: "Python unit tests for Helm chart"
+ name: "Unit tests Helm: ${{matrix.helm-test-package}}"
  runs-on: "${{needs.build-info.outputs.runs-on}}"
  needs: [build-info, wait-for-ci-images]
+ strategy:
+ fail-fast: false
+ matrix:
+ helm-test-package: ${{fromJson(needs.build-info.outputs.helm-test-packages)}}
  env:
  RUNS_ON: "${{needs.build-info.outputs.runs-on}}"
  PARALLEL_TEST_TYPES: "Helm"
@@ -834,8 +839,8 @@ jobs:
  - name: >
  Prepare breeze & CI image: ${{needs.build-info.outputs.default-python-version}}:${{env.IMAGE_TAG}}
  uses: ./.github/actions/prepare_breeze_and_image
- - name: "Helm Unit Tests"
- run: breeze testing helm-tests
+ - name: "Helm Unit Tests: ${{ matrix.helm-test-package }}"
+ run: breeze testing helm-tests --helm-test-package "${{ matrix.helm-test-package }}"
  - name: "Post Helm Tests"
  uses: ./.github/actions/post_tests
 

Dockerfile.ci

@@ -868,8 +868,13 @@ if [[ "${RUN_TESTS}" != "true" ]]; then
 fi
 set -u
 
-export RESULT_LOG_FILE="/files/test_result-${TEST_TYPE/\[*\]/}-${BACKEND}.xml"
-export WARNINGS_FILE="/files/warnings-${TEST_TYPE/\[*\]/}-${BACKEND}.txt"
+if [[ ${HELM_TEST_PACKAGE=} != "" ]]; then
+ export RESULT_LOG_FILE="/files/test_result-${TEST_TYPE/\[*\]/}-${HELM_TEST_PACKAGE}-${BACKEND}.xml"
+ export WARNINGS_FILE="/files/warnings-${TEST_TYPE/\[*\]/}-${HELM_TEST_PACKAGE}-${BACKEND}.txt"
+else
+ export RESULT_LOG_FILE="/files/test_result-${TEST_TYPE/\[*\]/}-${BACKEND}.xml"
+ export WARNINGS_FILE="/files/warnings-${TEST_TYPE/\[*\]/}-${BACKEND}.txt"
+fi
 
 EXTRA_PYTEST_ARGS=(
  "--verbosity=0"
@@ -1025,7 +1030,11 @@ else
  elif [[ ${TEST_TYPE:=""} == "WWW" ]]; then
  SELECTED_TESTS=("${WWW_TESTS[@]}")
  elif [[ ${TEST_TYPE:=""} == "Helm" ]]; then
- SELECTED_TESTS=("${HELM_CHART_TESTS[@]}")
+ if [[ ${HELM_TEST_PACKAGE=} != "" ]]; then
+ SELECTED_TESTS=("tests/charts/${HELM_TEST_PACKAGE}")
+ else
+ SELECTED_TESTS=("${HELM_CHART_TESTS[@]}")
+ fi
  elif [[ ${TEST_TYPE:=""} == "Integration" ]]; then
  if [[ ${SKIP_PROVIDER_TESTS:=""} == "true" ]]; then
  SELECTED_TESTS=("${NO_PROVIDERS_INTEGRATION_TESTS[@]}")

TESTING.rst

@@ -606,7 +606,7 @@ Example test here:
 
 .. code-block:: python
 
- from tests.charts.helm_template_generator import render_chart, render_k8s_object
+ from tests.charts.common.helm_template_generator import render_chart, render_k8s_object
 
  git_sync_basic = """
  dags:
@@ -634,6 +634,16 @@ following command (but it takes quite a long time even in a multi-processor mach
 
  breeze testing helm-tests
 
+You can also execute tests from a selected package only. Tests in ``tests/chart`` are grouped by packages
+so rather than running all tests, you can run only tests from a selected package. For example:
+
+.. code-block:: bash
+
+ breeze testing helm-tests --helm-test-package basic
+
+Will run all tests from ``tests/charts/basic`` package.
+
+
 You can also run Helm tests individually via the usual ``breeze`` command. Just enter breeze and run the
 tests with pytest as you would do with regular unit tests (you can add ``-n auto`` command to run Helm
 tests in parallel - unlike most of the regular unit tests of ours that require a database, the Helm tests are

dev/breeze/src/airflow_breeze/commands/testing_commands.py

@@ -25,7 +25,11 @@
 from click import IntRange
 
 from airflow_breeze.commands.ci_image_commands import rebuild_or_pull_ci_image_if_needed
-from airflow_breeze.global_constants import ALLOWED_TEST_TYPE_CHOICES, all_selective_test_types
+from airflow_breeze.global_constants import (
+ ALLOWED_HELM_TEST_PACKAGES,
+ ALLOWED_TEST_TYPE_CHOICES,
+ all_selective_test_types,
+)
 from airflow_breeze.params.build_prod_params import BuildProdParams
 from airflow_breeze.params.shell_params import ShellParams
 from airflow_breeze.utils.ci_group import ci_group
@@ -51,7 +55,7 @@
  option_verbose,
 )
 from airflow_breeze.utils.console import Output, get_console
-from airflow_breeze.utils.custom_param_types import NotVerifiedBetterChoice
+from airflow_breeze.utils.custom_param_types import BetterChoice, NotVerifiedBetterChoice
 from airflow_breeze.utils.docker_command_utils import (
  DOCKER_COMPOSE_COMMAND,
  get_env_variables_for_docker_commands,
@@ -545,11 +549,18 @@ def integration_tests(
 @option_github_repository
 @option_verbose
 @option_dry_run
+@click.option(
+ "--helm-test-package",
+ help="Package to tests",
+ default="all",
+ type=BetterChoice(ALLOWED_HELM_TEST_PACKAGES),
+)
 @click.argument("extra_pytest_args", nargs=-1, type=click.UNPROCESSED)
 def helm_tests(
  extra_pytest_args: tuple,
  image_tag: str | None,
  mount_sources: str,
+ helm_test_package: str,
  github_repository: str,
 ):
  exec_shell_params = ShellParams(
@@ -560,6 +571,8 @@ def helm_tests(
  env_variables = get_env_variables_for_docker_commands(exec_shell_params)
  env_variables["RUN_TESTS"] = "true"
  env_variables["TEST_TYPE"] = "Helm"
+ if helm_test_package != "all":
+ env_variables["HELM_TEST_PACKAGE"] = helm_test_package
  perform_environment_checks()
  cleanup_python_generated_files()
  cmd = [*DOCKER_COMPOSE_COMMAND, "run", "--service-ports", "--rm", "airflow"]

dev/breeze/src/airflow_breeze/commands/testing_commands_config.py

@@ -88,6 +88,7 @@
  "options": [
  "--image-tag",
  "--mount-sources",
+ "--helm-test-package",
  "--github-repository",
  ],
  },

dev/breeze/src/airflow_breeze/global_constants.py

@@ -115,6 +115,23 @@ class SelectiveUnitTestTypes(Enum):
  "Quarantine",
 ]
 
+
+@lru_cache(maxsize=None)
+def all_helm_test_packages() -> list[str]:
+ return sorted(
+ [
+ candidate.name
+ for candidate in (AIRFLOW_SOURCES_ROOT / "tests" / "charts").iterdir()
+ if candidate.is_dir()
+ ]
+ )
+
+
+ALLOWED_HELM_TEST_PACKAGES = [
+ "all",
+ *all_helm_test_packages(),
+]
+
 ALLOWED_PACKAGE_FORMATS = ["wheel", "sdist", "both"]
 ALLOWED_INSTALLATION_PACKAGE_FORMATS = ["wheel", "sdist"]
 ALLOWED_INSTALLATION_METHODS = [".", "apache-airflow"]

dev/breeze/src/airflow_breeze/utils/selective_checks.py

@@ -58,6 +58,7 @@
  KIND_VERSION,
  GithubEvents,
  SelectiveUnitTestTypes,
+ all_helm_test_packages,
  all_selective_test_types,
 )
 from airflow_breeze.utils.console import get_console
@@ -644,3 +645,7 @@ def debug_resources(self) -> bool:
  @cached_property
  def suspended_providers_folders(self) -> str:
  return " ".join(get_suspended_providers_folders())
+
+ @cached_property
+ def helm_test_packages(self) -> str:
+ return json.dumps(all_helm_test_packages())

docs/exts/docs_build/helm_chart_utils.py

@@ -17,10 +17,11 @@
 from __future__ import annotations
 
 import os
+from pathlib import Path
 
 import yaml
 
-CHART_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, os.pardir, "chart"))
+CHART_DIR = Path(__file__).resolve().parents[2] / "chart"
 CHART_YAML_PATH = os.path.join(CHART_DIR, "Chart.yaml")
 
 

images/breeze/output-commands-hash.txt

@@ -58,7 +58,7 @@ start-airflow:51fb2f709c7f5400a7e87c2fa06fafc4
 static-checks:543f0c776d0f198e80a0f75058445bb2
 stop:e5aa686b4e53707ced4039d8414d5cd6
 testing:docker-compose-tests:b86c044b24138af0659a05ed6331576c
-testing:helm-tests:94a442e7f3f63b34c4831a84d165690a
+testing:helm-tests:936cf28fd84ce4ff5113795fdae9624b
 testing:integration-tests:225ddb6243cce5fc64f4824b87adfd98
 testing:tests:86441445a2b521e8d5aee04d74978451
-testing:68efcf0731170e4ba2029121a5209e3a
+testing:2d95034763ee699f2e2fc1804f2fd7f0