Unit test bug fixes.

cosmos/__init__.py

@@ -10,8 +10,8 @@
 from cosmos.dataset import get_dbt_dataset
 
 # re-export the dag and task group
-from cosmos.dag import DbtDag
-from cosmos.task_group import DbtTaskGroup
+from cosmos.airflow.dag import DbtDag
+from cosmos.airflow.task_group import DbtTaskGroup
 from cosmos.operators.lazy_load import MissingPackage
 
 # re-export the operators

cosmos/dataset.py

@@ -1,21 +1,21 @@
-from typing import Any
+from typing import Any, Tuple
 
 
 try:
     from airflow.datasets import Dataset
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     from logging import getLogger
 
     logger = getLogger(__name__)
 
     class Dataset:  # type: ignore[no-redef]
         cosmos_override = True
 
-        def __init__(self, id: str, *args: tuple[Any], **kwargs: str):
+        def __init__(self, id: str, *args: Tuple[Any], **kwargs: str):
             self.id = id
             logger.warning("Datasets are not supported in Airflow < 2.5.0")
 
-        def __eq__(self, other: Dataset) -> bool:
+        def __eq__(self, other: "Dataset") -> bool:
             return self.id == other.id  # type: ignore[no-any-return]
 
 

cosmos/dbt/selector.py

@@ -2,7 +2,10 @@
 import logging
 from pathlib import Path
 
-from cosmos.dbt.graph import DbtNode
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cosmos.dbt.graph import DbtNode
 
 
 SUPPORTED_CONFIG = ["materialized", "schema", "tags"]

cosmos/operators/docker.py

@@ -153,7 +153,7 @@ class DbtRunOperationDockerOperator(DbtDockerBaseOperator):
     """
 
     ui_color = "#8194E0"
-    template_fields: Sequence[str] = "args"
+    template_fields: Sequence[str] = ("args",)
 
     def __init__(self, macro_name: str, args: Optional[dict[str, Any]] = None, **kwargs: str) -> None:
         self.macro_name = macro_name

cosmos/operators/kubernetes.py

@@ -161,7 +161,7 @@ class DbtRunOperationKubernetesOperator(DbtKubernetesBaseOperator):
     """
 
     ui_color = "#8194E0"
-    template_fields: Sequence[str] = "args"
+    template_fields: Sequence[str] = ("args",)
 
     def __init__(self, macro_name: str, args: Optional[dict[str, Any]] = None, **kwargs: str) -> None:
         self.macro_name = macro_name

cosmos/operators/local.py

@@ -6,7 +6,7 @@
 import signal
 import tempfile
 from pathlib import Path
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, Optional, Sequence, Tuple
 
 import yaml
 from airflow.compat.functools import cached_property
@@ -126,7 +126,7 @@ def store_compiled_sql(self, tmp_project_dir: str, context: Context, session: Se
         ).delete()
         session.add(rtif)
 
-    def run_subprocess(self, *args: tuple[Any], **kwargs: Any) -> FullOutputSubprocessResult:
+    def run_subprocess(self, *args: Tuple[Any], **kwargs: Any) -> FullOutputSubprocessResult:
         return self.subprocess_hook.run_command(*args, **kwargs)  # type: ignore[no-any-return]
 
     def get_profile_name(self, project_dir: str) -> str:
@@ -405,7 +405,7 @@ class DbtRunOperationLocalOperator(DbtLocalBaseOperator):
     """
 
     ui_color = "#8194E0"
-    template_fields: Sequence[str] = "args"
+    template_fields: Sequence[str] = ("args",)
 
     def __init__(self, macro_name: str, args: Optional[dict[str, Any]] = None, **kwargs: Any) -> None:
         self.macro_name = macro_name

cosmos/operators/virtualenv.py

@@ -3,7 +3,7 @@
 import logging
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Tuple
 
 
 from airflow.compat.functools import cached_property
@@ -85,9 +85,9 @@ def venv_dbt_path(
         self.log.info("Using dbt version %s available at %s", dbt_version, dbt_binary)
         return str(dbt_binary)
 
-    def run_subprocess(self, *args: tuple[Any], **kwargs: Any) -> FullOutputSubprocessResult:
-        command = kwargs["command"]
-
+    def run_subprocess(  # type: ignore[override]
+        self, *args: Tuple[Any], command: list[str], **kwargs: Any
+    ) -> FullOutputSubprocessResult:
         if self.py_requirements:
             command[0] = self.venv_dbt_path