docker-tests/tests/docker_tests/test_prod_image.py - airflow - Git at Google

 # 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

 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"

 try:
     from tomllib import loads as load_tomllib
 except ImportError:
     from tomli import loads as load_tomllib

 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", "sshtunnel"],
         "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]
             # FAB provider has import issues on Python 3.13
             if python_minor == "3.13":
                 import_names = [name for name in import_names if name != "airflow.providers.fab"]
         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"
	# 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

	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"

	try:
	from tomllib import loads as load_tomllib
	except ImportError:
	from tomli import loads as load_tomllib

	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", "sshtunnel"],
	"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]
	# FAB provider has import issues on Python 3.13
	if python_minor == "3.13":
	import_names = [name for name in import_names if name != "airflow.providers.fab"]
	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"