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
 import sys
 from importlib.util import find_spec
 from pathlib import Path

 import pytest
 from python_on_whales import DockerException

 from docker_tests.constants import SOURCE_ROOT
 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 = SOURCE_ROOT / "prod_image_installed_providers.txt"
 AIRFLOW_ROOT_PATH = Path(__file__).parents[2].resolve()

 sys.path.insert(0, AIRFLOW_ROOT_PATH.as_posix())  # make sure airflow root is imported
 from hatch_build import PRE_INSTALLED_PROVIDERS  # noqa: E402

 SLIM_IMAGE_PROVIDERS = [
     f"apache-airflow-providers-{provider_id.split('>=')[0].replace('.','-')}"
     for provider_id in PRE_INSTALLED_PROVIDERS
     if not provider_id.startswith("#")
 ]
 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("#")
 ]


 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 2 == ctx.value.return_code

     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 2 == ctx.value.return_code

     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 "2." 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 os.environ.get("TEST_SLIM_IMAGE"):
             packages_to_install = set(SLIM_IMAGE_PROVIDERS)
             package_file = "hatch_build.py"
         else:
             packages_to_install = set(REGULAR_IMAGE_PROVIDERS)
             package_file = PROD_IMAGE_PROVIDERS_FILE_PATH
         assert len(packages_to_install) != 0
         output = run_bash_in_docker("airflow providers list --output json", image=default_docker_image)
         providers = json.loads(output)
         packages_installed = set(d["package_name"] for d in providers)
         assert len(packages_installed) != 0

         assert (
             packages_to_install == packages_installed
         ), f"List of expected installed packages and image content mismatch. Check {package_file} file."

     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

     PACKAGE_IMPORTS = {
         "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"],
         "virtualenv": ["virtualenv"],
     }

     @pytest.mark.skipif(os.environ.get("TEST_SLIM_IMAGE") == "true", reason="Skipped with slim image")
     @pytest.mark.parametrize("package_name,import_names", PACKAGE_IMPORTS.items())
     def test_check_dependencies_imports(self, package_name, import_names, default_docker_image):
         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
	import sys
	from importlib.util import find_spec
	from pathlib import Path

	import pytest
	from python_on_whales import DockerException

	from docker_tests.constants import SOURCE_ROOT
	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 = SOURCE_ROOT / "prod_image_installed_providers.txt"
	AIRFLOW_ROOT_PATH = Path(__file__).parents[2].resolve()

	sys.path.insert(0, AIRFLOW_ROOT_PATH.as_posix()) # make sure airflow root is imported
	from hatch_build import PRE_INSTALLED_PROVIDERS # noqa: E402

	SLIM_IMAGE_PROVIDERS = [
	f"apache-airflow-providers-{provider_id.split('>=')[0].replace('.','-')}"
	for provider_id in PRE_INSTALLED_PROVIDERS
	if not provider_id.startswith("#")
	]
	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("#")
	]


	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 2 == ctx.value.return_code

	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 2 == ctx.value.return_code

	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 "2." 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 os.environ.get("TEST_SLIM_IMAGE"):
	packages_to_install = set(SLIM_IMAGE_PROVIDERS)
	package_file = "hatch_build.py"
	else:
	packages_to_install = set(REGULAR_IMAGE_PROVIDERS)
	package_file = PROD_IMAGE_PROVIDERS_FILE_PATH
	assert len(packages_to_install) != 0
	output = run_bash_in_docker("airflow providers list --output json", image=default_docker_image)
	providers = json.loads(output)
	packages_installed = set(d["package_name"] for d in providers)
	assert len(packages_installed) != 0

	assert (
	packages_to_install == packages_installed
	), f"List of expected installed packages and image content mismatch. Check {package_file} file."

	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

	PACKAGE_IMPORTS = {
	"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"],
	"virtualenv": ["virtualenv"],
	}

	@pytest.mark.skipif(os.environ.get("TEST_SLIM_IMAGE") == "true", reason="Skipped with slim image")
	@pytest.mark.parametrize("package_name,import_names", PACKAGE_IMPORTS.items())
	def test_check_dependencies_imports(self, package_name, import_names, default_docker_image):
	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"