diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 475090f5cb..d14ae5de72 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -32,7 +32,7 @@ jobs: architecture: "x64" - run: pip3 install hatch - - run: hatch run tests.py3.9-2.8:type-check + - run: hatch run tests.py3.9-2.8-1.9:type-check Run-Unit-Tests: runs-on: ubuntu-latest @@ -40,6 +40,7 @@ jobs: matrix: python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10"] + dbt-version: ["1.9"] exclude: - python-version: "3.11" airflow-version: "2.4" @@ -75,7 +76,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: unit-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: unit-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -86,16 +87,16 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - - name: Test Cosmos against Airflow ${{ matrix.airflow-version }} and Python ${{ matrix.python-version }} + - name: Test Cosmos against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt ${{ matrix.dbt-version }} run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-cov + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-cov - name: Upload coverage to Github uses: actions/upload-artifact@v4 with: - name: coverage-unit-test-${{ matrix.python-version }}-${{ matrix.airflow-version }} + name: coverage-unit-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} path: .coverage include-hidden-files: true @@ -106,6 +107,7 @@ jobs: matrix: python-version: ["3.8", "3.9", "3.10", "3.11"] airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10"] + dbt-version: ["1.9"] exclude: - python-version: "3.11" airflow-version: "2.4" @@ -139,7 +141,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: integration-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: integration-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}--${{ matrix.dbt-version }}${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -150,12 +152,12 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - - name: Test Cosmos against Airflow ${{ matrix.airflow-version }} and Python ${{ matrix.python-version }} + - name: Test Cosmos against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt ${{ matrix.dbt-version }} run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-setup - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration-setup + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration env: AIRFLOW__COSMOS__ENABLE_CACHE_DBT_LS: 0 AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ @@ -182,7 +184,7 @@ jobs: - name: Upload coverage to Github uses: actions/upload-artifact@v4 with: - name: coverage-integration-test-${{ matrix.python-version }}-${{ matrix.airflow-version }} + name: coverage-integration-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} path: .coverage include-hidden-files: true @@ -193,6 +195,7 @@ jobs: matrix: python-version: ["3.11"] airflow-version: ["2.6"] + dbt-version: ["1.9"] services: postgres: @@ -216,7 +219,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: integration-expensive-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: integration-expensive-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -227,12 +230,12 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - - name: Test Cosmos against Airflow ${{ matrix.airflow-version }} and Python ${{ matrix.python-version }} + - name: Test Cosmos against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt ${{ matrix.dbt-version }} run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-setup - DATABRICKS_UNIQUE_ID="${{github.run_id}}" hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-expensive + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration-setup + DATABRICKS_UNIQUE_ID="${{github.run_id}}" hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration-expensive env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres @@ -256,7 +259,7 @@ jobs: - name: Upload coverage to Github uses: actions/upload-artifact@v4 with: - name: coverage-integration-expensive-test-${{ matrix.python-version }}-${{ matrix.airflow-version }} + name: coverage-integration-expensive-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} path: .coverage include-hidden-files: true @@ -274,6 +277,7 @@ jobs: matrix: python-version: [ "3.11" ] airflow-version: [ "2.8" ] + dbt-version: [ "1.5" ] services: postgres: image: postgres @@ -296,7 +300,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: integration-dbt-1-5-4-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: integration-dbt-1-5-4-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -307,11 +311,11 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - name: Test Cosmos against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt 1.5.4 run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-dbt-1-5-4 + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration-dbt-1-5-4 env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres @@ -335,10 +339,68 @@ jobs: AIRFLOW__COSMOS__REMOTE_TARGET_PATH: "s3://cosmos-remote-cache/target_compiled/" AIRFLOW__COSMOS__REMOTE_TARGET_PATH_CONN_ID: aws_s3_conn + Run-Integration-Tests-DBT-Async: + needs: Authorize + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: [ "3.11" ] + airflow-version: [ "2.10" ] + dbt-version: ["1.5", "1.6", "1.7", "1.8", "1.9"] + services: + postgres: + image: postgres + env: + POSTGRES_PASSWORD: postgres + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + ports: + - 5432:5432 + + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha || github.ref }} + - uses: actions/cache@v4 + with: + path: | + ~/.cache/pip + .local/share/hatch/ + key: integration-dbt-async-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install packages and dependencies + run: | + python -m pip install uv + uv pip install --system hatch + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze + + - name: Test Cosmos against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt ${{ matrix.dbt-version }} + run: | + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-integration-dbt-async + env: + AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ + AIRFLOW_CONN_AWS_S3_CONN: ${{ secrets.AIRFLOW_CONN_AWS_S3_CONN }} + AIRFLOW_CONN_GCP_GS_CONN: ${{ secrets.AIRFLOW_CONN_GCP_GS_CONN }} + AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 90.0 + PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH + AIRFLOW__COSMOS__ENABLE_CACHE: 0 + AIRFLOW__COSMOS__REMOTE_TARGET_PATH: "s3://cosmos-remote-cache/target_compiled/" + AIRFLOW__COSMOS__REMOTE_TARGET_PATH_CONN_ID: aws_s3_conn + DBT_ADAPTER_VERSION: ${{ matrix.dbt-version }} + - name: Upload coverage to Github uses: actions/upload-artifact@v4 with: - name: coverage-integration-dbt-1-5-4-test-${{ matrix.python-version }}-${{ matrix.airflow-version }} + name: coverage-integration-dbt-async-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} path: .coverage include-hidden-files: true @@ -354,6 +416,7 @@ jobs: matrix: python-version: ["3.11"] airflow-version: ["2.8"] + dbt-version: ["1.9"] num-models: [1, 10, 50, 100, 500] services: postgres: @@ -376,7 +439,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: perf-test-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: perf-test-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -387,13 +450,13 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - - name: Run performance tests against against Airflow ${{ matrix.airflow-version }} and Python ${{ matrix.python-version }} + - name: Run performance tests against against Airflow ${{ matrix.airflow-version }}, Python ${{ matrix.python-version }} and dbt ${{ matrix.dbt-version }} id: run-performance-tests run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-performance-setup - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-performance + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-performance-setup + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-performance # read the performance results and set them as an env var for the next step # format: NUM_MODELS={num_models}\nTIME={end - start}\n @@ -423,6 +486,7 @@ jobs: matrix: python-version: [ "3.12" ] airflow-version: [ "2.10" ] + dbt-version: [ "1.9" ] steps: - uses: actions/checkout@v4 with: @@ -432,7 +496,7 @@ jobs: path: | ~/.cache/pip .local/share/hatch/ - key: coverage-integration-kubernetes-test-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} + key: coverage-integration-kubernetes-test-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -446,12 +510,12 @@ jobs: run: | python -m pip install uv uv pip install --system hatch - hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze + hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} run pip freeze - name: Run kubernetes tests run: | - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-kubernetes-setup - hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-kubernetes + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-kubernetes-setup + hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }}:test-kubernetes env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres @@ -477,7 +541,7 @@ jobs: - name: Upload coverage to Github uses: actions/upload-artifact@v4 with: - name: coverage-integration-kubernetes-test-${{ matrix.python-version }}-${{ matrix.airflow-version }} + name: coverage-integration-kubernetes-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ matrix.dbt-version }} path: .coverage include-hidden-files: true diff --git a/dev/dags/simple_dag_async.py b/dev/dags/simple_dag_async.py index 23e9836cff..d9afba493d 100644 --- a/dev/dags/simple_dag_async.py +++ b/dev/dags/simple_dag_async.py @@ -3,11 +3,14 @@ from pathlib import Path from cosmos import DbtDag, ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig, RenderConfig +from cosmos.constants import TestBehavior from cosmos.profiles import GoogleCloudServiceAccountDictProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) +DBT_ADAPTER_VERSION = os.getenv("DBT_ADAPTER_VERSION", "1.9") + profile_config = ProfileConfig( profile_name="default", target_name="dev", @@ -26,12 +29,9 @@ profile_config=profile_config, execution_config=ExecutionConfig( execution_mode=ExecutionMode.AIRFLOW_ASYNC, - async_py_requirements=["dbt-bigquery"], - ), - render_config=RenderConfig( - select=["path:models"], - # test_behavior=TestBehavior.NONE + async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], ), + render_config=RenderConfig(select=["path:models"], test_behavior=TestBehavior.NONE), # normal dag parameters schedule_interval=None, start_date=datetime(2023, 1, 1), diff --git a/pyproject.toml b/pyproject.toml index 122e06da28..24a00c4962 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -162,6 +162,7 @@ pre-install-commands = ["sh scripts/test/pre-install-airflow.sh {matrix:airflow} [[tool.hatch.envs.tests.matrix]] python = ["3.8", "3.9", "3.10", "3.11", "3.12"] airflow = ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10"] +dbt = ["1.5", "1.6", "1.7", "1.8", "1.9"] [tool.hatch.envs.tests.overrides] matrix.airflow.dependencies = [ @@ -176,6 +177,7 @@ test-integration = 'sh scripts/test/integration.sh' test-kubernetes = "sh scripts/test/integration-kubernetes.sh" test-kubernetes-setup = "sh scripts/test/kubernetes-setup.sh" test-integration-dbt-1-5-4 = 'sh scripts/test/integration-dbt-1-5-4.sh' +test-integration-dbt-async = 'sh scripts/test/integration-dbt-async.sh {matrix:dbt}' test-integration-expensive = 'sh scripts/test/integration-expensive.sh' test-integration-setup = 'sh scripts/test/integration-setup.sh' test-performance = 'sh scripts/test/performance.sh' diff --git a/scripts/test/integration-dbt-async.sh b/scripts/test/integration-dbt-async.sh new file mode 100644 index 0000000000..b4ff9e9195 --- /dev/null +++ b/scripts/test/integration-dbt-async.sh @@ -0,0 +1,29 @@ +#!/bin/bash + +set -x +set -e + +DBT_VERSION="$1" +echo "DBT_VERSION:" +echo "$DBT_VERSION" + + +pip uninstall dbt-adapters dbt-common dbt-core dbt-extractor dbt-postgres dbt-semantic-interfaces -y +pip install "dbt-postgres==$DBT_VERSION" "dbt-databricks==$DBT_VERSION" "dbt-bigquery==$DBT_VERSION" +export SOURCE_RENDERING_BEHAVIOR=all +rm -rf airflow.*; \ +airflow db init; \ + +if [ "$DBT_VERSION" = "1.7" ]; then + # Otherwise, we will get the following error: + # stderr: MessageToJson() got an unexpected keyword argument 'including_default_value_fields' + echo "DBT version is 1.7 — Installing protobuf==4.25.6..." + pip install protobuf==4.25.6 +fi + +rm -rf dbt/jaffle_shop/dbt_packages; +pytest -vv \ + --cov=cosmos \ + --cov-report=term-missing \ + --cov-report=xml \ + "tests/test_async_example_dag.py::test_example_dag[simple_dag_async]" diff --git a/scripts/test/integration.sh b/scripts/test/integration.sh index 985c99df57..bf19b4248f 100644 --- a/scripts/test/integration.sh +++ b/scripts/test/integration.sh @@ -19,5 +19,6 @@ pytest -vv \ --cov-report=xml \ -m 'integration' \ --ignore=tests/perf \ + --ignore=tests/test_async_example_dag.py \ --ignore=tests/test_example_k8s_dags.py \ -k 'not ( example_cosmos_python_models or example_virtualenv or jaffle_shop_kubernetes)' diff --git a/scripts/test/performance.sh b/scripts/test/performance.sh index 2023026d3b..97941e86b7 100644 --- a/scripts/test/performance.sh +++ b/scripts/test/performance.sh @@ -2,5 +2,6 @@ pytest -vv \ -s \ -m 'perf' \ --ignore=tests/test_example_dags.py \ + --ignore=tests/test_async_example_dag.py \ --ignore=tests/test_example_dags_no_connections.py \ --ignore=tests/test_example_k8s_dags.py diff --git a/scripts/test/unit-cov.sh b/scripts/test/unit-cov.sh index 8d8fe3589a..50cd268c40 100644 --- a/scripts/test/unit-cov.sh +++ b/scripts/test/unit-cov.sh @@ -6,5 +6,6 @@ pytest \ -m "not (integration or perf)" \ --ignore=tests/perf \ --ignore=tests/test_example_dags.py \ + --ignore=tests/test_async_example_dag.py \ --ignore=tests/test_example_dags_no_connections.py \ --ignore=tests/test_example_k8s_dags.py diff --git a/scripts/test/unit.sh b/scripts/test/unit.sh index b80aab72ea..39069c423a 100644 --- a/scripts/test/unit.sh +++ b/scripts/test/unit.sh @@ -3,5 +3,6 @@ pytest \ -m "not (integration or perf)" \ --ignore=tests/perf \ --ignore=tests/test_example_dags.py \ + --ignore=tests/test_async_example_dag.py \ --ignore=tests/test_example_dags_no_connections.py \ --ignore=tests/test_example_k8s_dags.py diff --git a/tests/test_async_example_dag.py b/tests/test_async_example_dag.py new file mode 100644 index 0000000000..5d70af6a64 --- /dev/null +++ b/tests/test_async_example_dag.py @@ -0,0 +1,73 @@ +# We already have tests/test_example_dags.py, but it doesn’t run against multiple dbt versions in CI. +# Some dbt versions have shown parsing issues with certain example DAGs — something we may need to address over time. +# With PR #1535, the goal is to test the async example DAG across multiple dbt versions. To prevent the CI job from +# failing early due to unrelated DAG parsing errors, PR #1535 introduces this new test_async_example_dag.py file. +# This file replicates tests/test_example_dags.py but excludes all DAGs except simple_async_dag by adding them to +# .airflowignore. This ensures the CI job focuses solely on testing simple_async_dag over multiple dbt versions +# without being disrupted by other DAG parsing issues. + +from __future__ import annotations + +from pathlib import Path + +try: + from functools import cache +except ImportError: + from functools import lru_cache as cache + + +import airflow +import pytest +from airflow.models.dagbag import DagBag +from airflow.utils.db import create_default_connections +from airflow.utils.session import provide_session +from packaging.version import Version + +EXAMPLE_DAGS_DIR = Path(__file__).parent.parent / "dev/dags" +ALL_FILES_TO_IGNORE = [ + f.name for f in EXAMPLE_DAGS_DIR.iterdir() if f.is_file() and f.suffix == ".py" and f.name != "simple_dag_async.py" +] + +AIRFLOW_IGNORE_FILE = EXAMPLE_DAGS_DIR / ".airflowignore" +AIRFLOW_VERSION = Version(airflow.__version__) + + +@provide_session +def get_session(session=None): + create_default_connections(session) + return session + + +@pytest.fixture() +def session(): + return get_session() + + +@cache +def get_dag_bag() -> DagBag: + """Create a DagBag by adding the files that are not supported to .airflowignore""" + + with open(AIRFLOW_IGNORE_FILE, "w+") as file: + for dagfile in ALL_FILES_TO_IGNORE: + print(f"Adding {dagfile} to .airflowignore") + file.writelines([f"{dagfile}\n"]) + + print(".airflowignore contents: ") + print(AIRFLOW_IGNORE_FILE.read_text()) + db = DagBag(EXAMPLE_DAGS_DIR, include_examples=False) + assert db.dags + assert not db.import_errors + return db + + +def get_dag_ids() -> list[str]: + dag_bag = get_dag_bag() + return dag_bag.dag_ids + + +@pytest.mark.integration +@pytest.mark.parametrize("dag_id", get_dag_ids()) +def test_example_dag(session, dag_id: str): + dag_bag = get_dag_bag() + dag = dag_bag.get_dag(dag_id) + dag.test()