import shutil from pathlib import Path from tempfile import mkdtemp from typing import Optional import pytest import yaml from dbt.version import __version__ as dbt_version from dbt_project import PYTEST_XDIST_WORKER, SCHEMA_NAME_SUFFIX, DbtProject from elementary.clients.dbt.factory import RunnerMethod from env import Environment from logger import get_logger from packaging import version DBT_PROJECT_PATH = Path(__file__).parent.parent / "dbt_project" DBT_FUSION_SUPPORTED_TARGETS = [ "snowflake", "bigquery", "redshift", "databricks_catalog", ] logger = get_logger(__name__) def pytest_addoption(parser): parser.addoption("--target", action="store", default="postgres") parser.addoption("--skip-init", action="store_true", default=False) parser.addoption("--clear-on-end", action="store_true", default=False) parser.addoption("--runner-method", action="store", default=None) @pytest.fixture(scope="session") def project_dir_copy(runner_method: Optional[RunnerMethod]): dbt_project_copy_dir = mkdtemp(prefix="integration_tests_project_") try: shutil.copytree( DBT_PROJECT_PATH, dbt_project_copy_dir, dirs_exist_ok=True, symlinks=True, ) _edit_packages_yml_to_include_absolute_elementary_package_path( dbt_project_copy_dir ) _remove_python_models_for_dbt_fusion(dbt_project_copy_dir, runner_method) yield dbt_project_copy_dir finally: shutil.rmtree(dbt_project_copy_dir) def _edit_packages_yml_to_include_absolute_elementary_package_path( project_dir_copy: str, ): logger.info( f"Editing packages.yml to include absolute elementary package path for project {project_dir_copy}" ) packages_yml_path = Path(project_dir_copy) / "packages.yml" with packages_yml_path.open("r") as packages_yml_file: packages_yml = yaml.safe_load(packages_yml_file) packages_yml["packages"][0]["local"] = str( (DBT_PROJECT_PATH / packages_yml["packages"][0]["local"]).resolve() ) with packages_yml_path.open("w") as packages_yml_file: yaml.dump(packages_yml, packages_yml_file) def _remove_python_models_for_dbt_fusion( project_dir_copy: str, runner_method: Optional[RunnerMethod] ): if runner_method != RunnerMethod.FUSION: return logger.info(f"Removing python tests for project {project_dir_copy}") # walk on the models dir and delete python files for path in (Path(project_dir_copy) / "models").rglob("*.py"): path.unlink() @pytest.fixture(scope="session", autouse=True) def init_tests_env( target: str, skip_init: bool, clear_on_end: bool, project_dir_copy: str, runner_method: Optional[RunnerMethod], ): env = Environment(target, project_dir_copy, runner_method) if not skip_init: logger.info( "Initializing test environment (worker=%s, schema_suffix='%s')", PYTEST_XDIST_WORKER or "main", SCHEMA_NAME_SUFFIX, ) env.clear() env.init() logger.info( "Initialization complete (worker=%s)", PYTEST_XDIST_WORKER or "main", ) yield if clear_on_end: logger.info("Clearing tests environment") env.clear() logger.info("Clearing complete") @pytest.fixture(autouse=True) def skip_by_targets(request, target: str): if request.node.get_closest_marker("skip_targets"): skipped_targets = request.node.get_closest_marker("skip_targets").args[0] if target in skipped_targets: pytest.skip("Test unsupported for target: {}".format(target)) @pytest.fixture(autouse=True) def only_on_targets(request, target: str): if request.node.get_closest_marker("only_on_targets"): requested_targets = request.node.get_closest_marker("only_on_targets").args[0] if target not in requested_targets: pytest.skip("Test unsupported for target: {}".format(target)) @pytest.fixture(autouse=True) def skip_for_dbt_fusion(request, runner_method: Optional[RunnerMethod]): if request.node.get_closest_marker("skip_for_dbt_fusion"): if runner_method == RunnerMethod.FUSION: pytest.skip("Test unsupported for dbt fusion") @pytest.fixture(autouse=True) def requires_dbt_version(request): if request.node.get_closest_marker("requires_dbt_version"): required_version = request.node.get_closest_marker("requires_dbt_version").args[ 0 ] if version.parse(dbt_version) < version.parse(required_version): pytest.skip( "Test requires dbt version {} or above, but {} is installed.".format( required_version, dbt_version ) ) @pytest.fixture def dbt_project( target: str, project_dir_copy: str, runner_method: Optional[RunnerMethod] ) -> DbtProject: return DbtProject(target, project_dir_copy, runner_method) @pytest.fixture(scope="session") def target(request) -> str: return request.config.getoption("--target") @pytest.fixture(scope="session") def skip_init(request) -> bool: return request.config.getoption("--skip-init") @pytest.fixture(scope="session") def clear_on_end(request) -> bool: return request.config.getoption("--clear-on-end") @pytest.fixture(scope="session") def runner_method(request, target: str) -> Optional[RunnerMethod]: runner_method_str = request.config.getoption("--runner-method") if runner_method_str: runner_method = RunnerMethod(runner_method_str) if ( runner_method == RunnerMethod.FUSION and target not in DBT_FUSION_SUPPORTED_TARGETS ): raise ValueError(f"Fusion runner is not supported for target: {target}") return runner_method return None @pytest.fixture def test_id(request) -> str: if request.cls: return f"{request.cls.__name__}_{request.node.name}" return request.node.name