diff --git a/dev/dags/jaffle_shop_kubernetes.py b/dev/dags/jaffle_shop_kubernetes.py index 10877f5e6..a50a8e454 100644 --- a/dev/dags/jaffle_shop_kubernetes.py +++ b/dev/dags/jaffle_shop_kubernetes.py @@ -9,22 +9,21 @@ https://astronomer.github.io/astronomer-cosmos/getting_started/kubernetes.html#kubernetes """ + from airflow import DAG from airflow.providers.cncf.kubernetes.secret import Secret from pendulum import datetime from cosmos import ( - ProfileConfig, DbtSeedKubernetesOperator, + ProfileConfig, ) from cosmos.profiles import PostgresUserPasswordProfileMapping DBT_IMAGE = "dbt-jaffle-shop:1.0.0" -project_seeds = [ - {"project": "jaffle_shop", "seeds": ["raw_customers", "raw_payments", "raw_orders"]} -] +project_seeds = [{"project": "jaffle_shop", "seeds": ["raw_customers", "raw_payments", "raw_orders"]}] postgres_password_secret = Secret( deploy_type="env", @@ -61,7 +60,7 @@ conn_id="postgres_default", profile_args={ "schema": "public", - } - ) - ) + }, + ), + ), ) diff --git a/scripts/test/kubernetes-setup.sh b/scripts/test/kubernetes-setup.sh index d37ee71f0..bce7fb58a 100644 --- a/scripts/test/kubernetes-setup.sh +++ b/scripts/test/kubernetes-setup.sh @@ -74,4 +74,3 @@ kubectl get svc --namespace default echo "pg log" kubectl logs postgres-postgresql-0 -c postgresql kubectl describe pod postgres-postgresql-0 - diff --git a/tests/test_example_dags.py b/tests/test_example_dags.py index d25950124..e10093bb5 100644 --- a/tests/test_example_dags.py +++ b/tests/test_example_dags.py @@ -7,6 +7,7 @@ except ImportError: from functools import lru_cache as cache + import airflow import pytest from airflow.models.dagbag import DagBag @@ -31,6 +32,7 @@ IGNORED_DAG_FILES = ["performance_dag.py"] + # Sort descending based on Versions and convert string to an actual version MIN_VER_DAG_FILE_VER: dict[Version, list[str]] = { Version(version): MIN_VER_DAG_FILE[version] for version in sorted(MIN_VER_DAG_FILE, key=Version, reverse=True) @@ -49,7 +51,7 @@ def session(): @cache -def get_dag_bag(in_kube: bool = False) -> DagBag: +def get_dag_bag() -> DagBag: """Create a DagBag by adding the files that are not supported to .airflowignore""" if AIRFLOW_VERSION in PARTIALLY_SUPPORTED_AIRFLOW_VERSIONS: return DagBag(dag_folder=None, include_examples=False) @@ -70,9 +72,6 @@ def get_dag_bag(in_kube: bool = False) -> DagBag: if DBT_VERSION < Version("1.6.0"): file.writelines(["example_model_version.py\n"]) - if not in_kube: - file.writelines(["jaffle_shop_kubernetes.py\n"]) - print(".airflowignore contents: ") print(AIRFLOW_IGNORE_FILE.read_text()) db = DagBag(EXAMPLE_DAGS_DIR, include_examples=False) @@ -81,8 +80,8 @@ def get_dag_bag(in_kube: bool = False) -> DagBag: return db -def get_dag_ids(in_kube: bool = False) -> list[str]: - dag_bag = get_dag_bag(in_kube) +def get_dag_ids() -> list[str]: + dag_bag = get_dag_bag() return dag_bag.dag_ids diff --git a/tests/test_example_k8s_dags.py b/tests/test_example_k8s_dags.py index 213b03950..2b24fd466 100644 --- a/tests/test_example_k8s_dags.py +++ b/tests/test_example_k8s_dags.py @@ -3,11 +3,11 @@ import pytest from airflow.models.dagbag import DagBag - -from . import utils as test_utils from airflow.utils.db import create_default_connections from airflow.utils.session import provide_session +from . import utils as test_utils + EXAMPLE_DAGS_DIR = Path(__file__).parent.parent / "dev/dags" AIRFLOW_IGNORE_FILE = EXAMPLE_DAGS_DIR / ".airflowignore" @@ -39,6 +39,6 @@ def get_all_dag_files(): def test_example_dag_kubernetes(session): get_all_dag_files() db = DagBag(EXAMPLE_DAGS_DIR, include_examples=False) - #for dag_id in KUBERNETES_DAG_FILES: + # for dag_id in KUBERNETES_DAG_FILES: dag = db.get_dag("jaffle_shop_kubernetes") test_utils.run_dag(dag)