Use python-on-whales in docker tests

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

@@ -155,6 +155,7 @@
 )
 from airflow_breeze.utils.shared_options import get_dry_run, get_verbose
 from airflow_breeze.utils.versions import is_pre_release
+from airflow_breeze.utils.virtualenv_utils import create_pip_command, create_venv
 
 argument_provider_packages = click.argument(
     "provider_packages",
@@ -449,36 +450,8 @@ def _check_sdist_to_wheel_dists(dists_info: tuple[DistributionPackageInfo, ...])
                 continue
 
             if not venv_created:
-                venv_path = (Path(tmp_dir_name) / ".venv").resolve().absolute()
-                venv_command_result = run_command(
-                    [sys.executable, "-m", "venv", venv_path.as_posix()],
-                    check=False,
-                    capture_output=True,
-                )
-                if venv_command_result.returncode != 0:
-                    get_console().print(
-                        f"[error]Error when initializing virtualenv in {venv_path.as_posix()}:[/]\n"
-                        f"{venv_command_result.stdout}\n{venv_command_result.stderr}"
-                    )
-                python_path = venv_path / "bin" / "python"
-                if not python_path.exists():
-                    get_console().print(
-                        f"\n[errors]Python interpreter is not exist in path {python_path}. Exiting!\n"
-                    )
-                    sys.exit(1)
-                pip_command = (python_path.as_posix(), "-m", "pip")
-                result = run_command(
-                    [*pip_command, "install", f"pip=={AIRFLOW_PIP_VERSION}"],
-                    check=False,
-                    capture_output=True,
-                    text=True,
-                )
-                if result.returncode != 0:
-                    get_console().print(
-                        f"[error]Error when installing pip in {venv_path.as_posix()}[/]\n"
-                        f"{result.stdout}\n{result.stderr}"
-                    )
-                    sys.exit(1)
+                python_path = create_venv(Path(tmp_dir_name) / ".venv", pip_version=AIRFLOW_PIP_VERSION)
+                pip_command = create_pip_command(python_path)
                 venv_created = True
 
             returncode = _check_sdist_to_wheel(di, pip_command, str(tmp_dir_name))
@@ -492,7 +465,7 @@ def _check_sdist_to_wheel_dists(dists_info: tuple[DistributionPackageInfo, ...])
         sys.exit(1)
 
 
-def _check_sdist_to_wheel(dist_info: DistributionPackageInfo, pip_command: tuple[str, ...], cwd: str) -> int:
+def _check_sdist_to_wheel(dist_info: DistributionPackageInfo, pip_command: list[str], cwd: str) -> int:
     get_console().print(
         f"[info]Validate build wheel from sdist distribution for package {dist_info.package!r}.[/]"
     )

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

@@ -22,10 +22,15 @@
 from itertools import chain
 from subprocess import DEVNULL
 
+from airflow_breeze.global_constants import PIP_VERSION
 from airflow_breeze.utils.console import Output, get_console
 from airflow_breeze.utils.packages import get_excluded_provider_folders, get_suspended_provider_folders
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT
 from airflow_breeze.utils.run_utils import run_command
+from airflow_breeze.utils.virtualenv_utils import create_temp_venv
+
+DOCKER_TESTS_ROOT = AIRFLOW_SOURCES_ROOT / "docker_tests"
+DOCKER_TESTS_REQUIREMENTS = DOCKER_TESTS_ROOT / "requirements.txt"
 
 
 def verify_an_image(
@@ -47,19 +52,20 @@ def verify_an_image(
         return command_result.returncode, f"Testing {image_type} python {image_name}"
     pytest_args = ("-n", str(os.cpu_count()), "--color=yes")
     if image_type == "PROD":
-        test_path = AIRFLOW_SOURCES_ROOT / "docker_tests" / "test_prod_image.py"
+        test_path = DOCKER_TESTS_ROOT / "test_prod_image.py"
     else:
-        test_path = AIRFLOW_SOURCES_ROOT / "docker_tests" / "test_ci_image.py"
+        test_path = DOCKER_TESTS_ROOT / "test_ci_image.py"
     env = os.environ.copy()
     env["DOCKER_IMAGE"] = image_name
     if slim_image:
         env["TEST_SLIM_IMAGE"] = "true"
-    command_result = run_command(
-        [sys.executable, "-m", "pytest", str(test_path), *pytest_args, *extra_pytest_args],
-        env=env,
-        output=output,
-        check=False,
-    )
+    with create_temp_venv(pip_version=PIP_VERSION, requirements_file=DOCKER_TESTS_REQUIREMENTS) as py_exe:
+        command_result = run_command(
+            [py_exe, "-m", "pytest", str(test_path), *pytest_args, *extra_pytest_args],
+            env=env,
+            output=output,
+            check=False,
+        )
     return command_result.returncode, f"Testing {image_type} python {image_name}"
 
 
@@ -73,16 +79,17 @@ def run_docker_compose_tests(
         get_console().print(f"[error]Error when inspecting PROD image: {command_result.returncode}[/]")
         return command_result.returncode, f"Testing docker-compose python with {image_name}"
     pytest_args = ("--color=yes",)
-    test_path = AIRFLOW_SOURCES_ROOT / "docker_tests" / "test_docker_compose_quick_start.py"
+    test_path = DOCKER_TESTS_ROOT / "test_docker_compose_quick_start.py"
     env = os.environ.copy()
     env["DOCKER_IMAGE"] = image_name
     if skip_docker_compose_deletion:
         env["SKIP_DOCKER_COMPOSE_DELETION"] = "true"
-    command_result = run_command(
-        [sys.executable, "-m", "pytest", str(test_path), *pytest_args, *extra_pytest_args],
-        env=env,
-        check=False,
-    )
+    with create_temp_venv(pip_version=PIP_VERSION, requirements_file=DOCKER_TESTS_REQUIREMENTS) as py_exe:
+        command_result = run_command(
+            [py_exe, "-m", "pytest", str(test_path), *pytest_args, *extra_pytest_args],
+            env=env,
+            check=False,
+        )
     return command_result.returncode, f"Testing docker-compose python with {image_name}"
 
 

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

@@ -0,0 +1,100 @@
+# 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 contextlib
+import sys
+import tempfile
+from pathlib import Path
+from typing import Generator
+
+from airflow_breeze.utils.console import get_console
+from airflow_breeze.utils.run_utils import run_command
+
+
+def create_pip_command(python: str | Path) -> list[str]:
+    return [python.as_posix() if hasattr(python, "as_posix") else str(python), "-m", "pip"]
+
+
+def create_venv(
+    venv_path: str | Path,
+    python: str | None = None,
+    pip_version: str | None = None,
+    requirements_file: str | Path | None = None,
+) -> str:
+    venv_path = Path(venv_path).resolve().absolute()
+    venv_command_result = run_command(
+        [python or sys.executable, "-m", "venv", venv_path.as_posix()],
+        check=False,
+        capture_output=True,
+    )
+    if venv_command_result.returncode != 0:
+        get_console().print(
+            f"[error]Error when initializing virtualenv in {venv_path.as_posix()}:[/]\n"
+            f"{venv_command_result.stdout}\n{venv_command_result.stderr}"
+        )
+        sys.exit(venv_command_result.returncode)
+    python_path = venv_path / "bin" / "python"
+    if not python_path.exists():
+        get_console().print(f"\n[errors]Python interpreter is not exist in path {python_path}. Exiting!\n")
+        sys.exit(1)
+    pip_command = create_pip_command(python_path)
+    if pip_version:
+        result = run_command(
+            [*pip_command, "install", f"pip=={pip_version}", "-q"],
+            check=False,
+            capture_output=False,
+            text=True,
+        )
+        if result.returncode != 0:
+            get_console().print(
+                f"[error]Error when installing pip in {venv_path.as_posix()}[/]\n"
+                f"{result.stdout}\n{result.stderr}"
+            )
+            sys.exit(result.returncode)
+    if requirements_file:
+        requirements_file = Path(requirements_file).absolute().as_posix()
+        result = run_command(
+            [*pip_command, "install", "-r", requirements_file, "-q"],
+            check=True,
+            capture_output=False,
+            text=True,
+        )
+        if result.returncode != 0:
+            get_console().print(
+                f"[error]Error when installing packages from {requirements_file}[/]\n"
+                f"{result.stdout}\n{result.stderr}"
+            )
+            sys.exit(result.returncode)
+    return python_path.as_posix()
+
+
+@contextlib.contextmanager
+def create_temp_venv(
+    python: str | None = None,
+    pip_version: str | None = None,
+    requirements_file: str | Path | None = None,
+    prefix: str | None = None,
+) -> Generator[str, None, None]:
+    with tempfile.TemporaryDirectory(prefix=prefix) as tmp_dir_name:
+        yield create_venv(
+            Path(tmp_dir_name) / ".venv",
+            python=python,
+            pip_version=pip_version,
+            requirements_file=requirements_file,
+        )

docker_tests/conftest.py

@@ -0,0 +1,29 @@
+# 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 os
+
+import pytest
+
+from docker_tests.constants import DEFAULT_DOCKER_IMAGE
+
+
+@pytest.fixture
+def default_docker_image() -> str:
+    return os.environ.get("DOCKER_IMAGE") or DEFAULT_DOCKER_IMAGE

docker_tests/constants.py

@@ -16,6 +16,11 @@
 # under the License.
 from __future__ import annotations
 
+import os
 from pathlib import Path
 
 SOURCE_ROOT = Path(__file__).resolve().parents[1]
+
+DEFAULT_PYTHON_MAJOR_MINOR_VERSION = "3.8"
+DEFAULT_DOCKER_IMAGE = f"ghcr.io/apache/airflow/main/prod/python{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}:latest"
+DOCKER_IMAGE = os.environ.get("DOCKER_IMAGE") or DEFAULT_DOCKER_IMAGE