airflow/docker-tests/tests/docker_tests/test_prod_image.py

# 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 json
import os
from importlib.util import find_spec
from tomllib import loads as load_tomllib

import pytest
import yaml
from python_on_whales import DockerException

from docker_tests.constants import AIRFLOW_ROOT_PATH
from docker_tests.docker_utils import (
    display_dependency_conflict_message,
    run_airflow_cmd_in_docker,
    run_bash_in_docker,
    run_cmd_in_docker,
    run_python_in_docker,
)

PROD_IMAGE_PROVIDERS_FILE_PATH = AIRFLOW_ROOT_PATH / "prod_image_installed_providers.txt"
AIRFLOW_CORE_PYPROJECT_TOML = AIRFLOW_ROOT_PATH / "airflow-core" / "pyproject.toml"

airflow_core_pyproject_toml = load_tomllib(AIRFLOW_CORE_PYPROJECT_TOML.read_text())

SLIM_IMAGE_PROVIDERS = []

for dependency in airflow_core_pyproject_toml["project"]["dependencies"]:
    if dependency.startswith("apache-airflow-providers"):
        SLIM_IMAGE_PROVIDERS.append(dependency.split(">=")[0])

REGULAR_IMAGE_PROVIDERS = [
    f"apache-airflow-providers-{provider_id.split('>=')[0].replace('.', '-')}"
    for provider_id in PROD_IMAGE_PROVIDERS_FILE_PATH.read_text().splitlines()
    if not provider_id.startswith("#")
]

testing_slim_image = os.environ.get("TEST_SLIM_IMAGE", str(False)).lower() in ("true", "1", "yes")


def _get_provider_python_exclusions() -> dict[str, list[str]]:
    """Return mapping of provider_id -> list of excluded Python minor versions from provider.yaml."""
    exclusions: dict[str, list[str]] = {}
    providers_root = AIRFLOW_ROOT_PATH / "providers"
    for line in PROD_IMAGE_PROVIDERS_FILE_PATH.read_text().splitlines():
        provider_id = line.split(">=")[0].strip()
        if not provider_id or provider_id.startswith("#"):
            continue
        provider_yaml_path = providers_root / provider_id.replace(".", "/") / "provider.yaml"
        if not provider_yaml_path.exists():
            continue
        with open(provider_yaml_path) as f:
            data = yaml.safe_load(f)
        excluded = data.get("excluded-python-versions", [])
        if excluded:
            exclusions[provider_id] = [str(v) for v in excluded]
    return exclusions


PROVIDER_PYTHON_EXCLUSIONS = _get_provider_python_exclusions()


def _get_python_minor_version(default_docker_image: str) -> str:
    """Get Python minor version (e.g. '3.14') from the Docker image."""
    python_version = run_bash_in_docker("python --version", image=default_docker_image)
    return ".".join(python_version.strip().split()[1].split(".")[:2])


def _get_excluded_provider_ids(python_minor: str) -> set[str]:
    """Return set of provider IDs excluded for the given Python minor version."""
    return {pid for pid, versions in PROVIDER_PYTHON_EXCLUSIONS.items() if python_minor in versions}


class TestCommands:
    def test_without_command(self, default_docker_image):
        """Checking the image without a command. It should return non-zero exit code."""
        with pytest.raises(DockerException) as ctx:
            run_cmd_in_docker(image=default_docker_image)
        assert ctx.value.return_code == 2

    def test_airflow_command(self, default_docker_image):
        """Checking 'airflow' command. It should return non-zero exit code."""
        with pytest.raises(DockerException) as ctx:
            run_airflow_cmd_in_docker(image=default_docker_image)
        assert ctx.value.return_code == 2

    def test_airflow_version(self, default_docker_image):
        """Checking 'airflow version' command. It should return zero exit code."""
        output = run_airflow_cmd_in_docker(["version"], image=default_docker_image)
        assert "3." in output

    def test_python_version(self, default_docker_image):
        """Checking 'python --version' command. It should return zero exit code."""
        output = run_cmd_in_docker(cmd=["python", "--version"], image=default_docker_image)
        assert "Python 3." in output

    def test_bash_version(self, default_docker_image):
        """Checking 'bash --version' command  It should return zero exit code."""
        output = run_cmd_in_docker(cmd=["bash", "--version"], image=default_docker_image)
        assert "GNU bash," in output


class TestPythonPackages:
    def test_required_providers_are_installed(self, default_docker_image):
        if testing_slim_image:
            packages_to_install = set(SLIM_IMAGE_PROVIDERS)
        else:
            python_minor = _get_python_minor_version(default_docker_image)
            excluded_ids = _get_excluded_provider_ids(python_minor)
            excluded_packages = {f"apache-airflow-providers-{pid.replace('.', '-')}" for pid in excluded_ids}
            packages_to_install = set(REGULAR_IMAGE_PROVIDERS) - excluded_packages
        assert len(packages_to_install) != 0
        output = run_bash_in_docker(
            "airflow providers list --output json",
            image=default_docker_image,
            envs={"AIRFLOW__LOGGING__LOGGING_LEVEL": "ERROR"},
        )
        providers = json.loads(output)
        packages_installed = set(d["package_name"] for d in providers)
        assert len(packages_installed) != 0

        message = "List of expected installed packages and image content mismatch.\n"
        message += f"Expected packages: {sorted(packages_to_install)}\n"
        message += f"Installed packages: {sorted(packages_installed)}\n"
        if testing_slim_image:
            message += (
                f"Please check the {AIRFLOW_CORE_PYPROJECT_TOML} file for the "
                f"'apache-airflow-providers-' packages in `dependencies`.\n"
            )
        else:
            message += f"Please check the {PROD_IMAGE_PROVIDERS_FILE_PATH} file for the expected packages.\n"

        assert packages_to_install == packages_installed, message

    def test_pip_dependencies_conflict(self, default_docker_image):
        try:
            run_bash_in_docker("pip check", image=default_docker_image)
        except DockerException:
            display_dependency_conflict_message()
            raise

    REGULAR_PACKAGE_IMPORTS: dict[str, list[str]] = {
        "amazon": ["boto3", "botocore", "watchtower"],
        "async": ["gevent", "eventlet", "greenlet"],
        "azure": [
            "azure.batch",
            "azure.cosmos",
            "azure.datalake.store",
            "azure.identity",
            "azure.keyvault.secrets",
            "azure.kusto.data",
            "azure.mgmt.containerinstance",
            "azure.mgmt.datalake.store",
            "azure.mgmt.resource",
            "azure.storage",
        ],
        "celery": ["celery", "flower", "vine"],
        "cncf.kubernetes": ["kubernetes", "cryptography"],
        "docker": ["docker"],
        "elasticsearch": ["elasticsearch"],
        "google": [
            "OpenSSL",
            # "google.ads", Remove google ads as it is vendored in google provider now
            "googleapiclient",
            "google.auth",
            "google_auth_httplib2",
            "google.cloud.automl",
            "google.cloud.bigquery_datatransfer",
            "google.cloud.bigtable",
            "google.cloud.container",
            "google.cloud.datacatalog",
            "google.cloud.dataproc",
            "google.cloud.dlp",
            "google.cloud.kms",
            "google.cloud.language",
            "google.cloud.logging",
            "google.cloud.memcache",
            "google.cloud.monitoring",
            "google.cloud.oslogin",
            "google.cloud.pubsub",
            "google.cloud.redis",
            "google.cloud.secretmanager",
            "google.cloud.spanner",
            "google.cloud.speech",
            "google.cloud.storage",
            "google.cloud.tasks",
            "google.cloud.texttospeech",
            "google.cloud.translate",
            "google.cloud.videointelligence",
            "google.cloud.vision",
        ],
        "grpc": ["grpc", "google.auth", "google_auth_httplib2"],
        "hashicorp": ["hvac"],
        "ldap": ["ldap"],
        "mysql": ["MySQLdb", *(["mysql"] if bool(find_spec("mysql")) else [])],
        "postgres": ["psycopg2"],
        "pyodbc": ["pyodbc"],
        "redis": ["redis"],
        "sendgrid": ["sendgrid"],
        "sftp/ssh": ["paramiko"],
        "slack": ["slack_sdk"],
        "statsd": ["statsd"],
        "providers": [provider[len("apache-") :].replace("-", ".") for provider in REGULAR_IMAGE_PROVIDERS],
    }

    SLIM_PACKAGE_IMPORTS: dict[str, list[str]] = {
        "providers": [provider[len("apache-") :].replace("-", ".") for provider in SLIM_IMAGE_PROVIDERS]
    }

    @pytest.mark.parametrize(
        ("package_name", "import_names"),
        SLIM_PACKAGE_IMPORTS.items() if testing_slim_image else REGULAR_PACKAGE_IMPORTS.items(),
    )
    def test_check_dependencies_imports(
        self, package_name: str, import_names: list[str], default_docker_image: str
    ):
        python_minor = _get_python_minor_version(default_docker_image)
        excluded_ids = _get_excluded_provider_ids(python_minor)
        # Skip individual provider test cases if the provider is excluded for this Python version
        if package_name in excluded_ids:
            pytest.skip(f"Provider {package_name} is excluded for Python {python_minor}")
        if package_name == "providers":
            excluded_imports = {f"airflow.providers.{pid}" for pid in excluded_ids}
            import_names = [name for name in import_names if name not in excluded_imports]
        if import_names:
            run_python_in_docker(f"import {','.join(import_names)}", image=default_docker_image)

    def test_there_is_no_opt_airflow_airflow_folder(self, default_docker_image):
        output = run_bash_in_docker(
            "find /opt/airflow/airflow/ 2>/dev/null | wc -l", image=default_docker_image
        )
        assert output == "0"


class TestExecuteAsRoot:
    def test_execute_airflow_as_root(self, default_docker_image):
        run_cmd_in_docker(
            cmd=["airflow", "info"],
            user=0,
            envs={"PYTHONDONTWRITEBYTECODE": "true"},
            image=default_docker_image,
        )

    def test_run_custom_python_packages_as_root(self, tmp_path, default_docker_image):
        (tmp_path / "__init__.py").write_text("")
        (tmp_path / "awesome.py").write_text('print("Awesome")')

        output = run_cmd_in_docker(
            envs={"PYTHONPATH": "/custom/mount", "PYTHONDONTWRITEBYTECODE": "true"},
            volumes=[(tmp_path.as_posix(), "/custom/mount")],
            user=0,
            cmd=["python", "-c", "import awesome"],
            image=default_docker_image,
        )
        assert output.strip() == "Awesome"