Skip to content

Commit

Permalink
Add test to confirm behaviour of converter on cache dir
Browse files Browse the repository at this point in the history
  • Loading branch information
tatiana committed Apr 24, 2024
1 parent 0a64de8 commit e3fdce9
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
1 change: 1 addition & 0 deletions cosmos/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from cosmos.constants import DEFAULT_COSMOS_CACHE_DIR_NAME

# In MacOS users may want to set the envvar `TMPDIR` if they do not want the value of the temp directory to change
DEFAULT_CACHE_DIR = Path(tempfile.gettempdir(), DEFAULT_COSMOS_CACHE_DIR_NAME)
cache_dir = Path(conf.get("cosmos", "cache_dir", fallback=DEFAULT_CACHE_DIR) or DEFAULT_CACHE_DIR)
enable_cache = conf.get("cosmos", "enable_cache", fallback=True) or True
70 changes: 70 additions & 0 deletions tests/test_converter.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import tempfile
from datetime import datetime
from pathlib import Path
from unittest.mock import MagicMock, patch
Expand Down Expand Up @@ -472,6 +473,75 @@ def test_converter_invocation_mode_added_to_task_args(
assert "invocation_mode" not in kwargs["task_args"]


@patch("cosmos.config.ProjectConfig.validate_project")
@patch("cosmos.converter.validate_initial_user_config")
@patch("cosmos.converter.DbtGraph")
@patch("cosmos.converter.build_airflow_graph")
def test_converter_uses_cache_dir(
mock_build_airflow_graph,
mock_dbt_graph,
mock_user_config,
mock_validate_project,
):
"""Tests that DbtGraph and operator and Airflow task args contain expected cache dir ."""
project_config = ProjectConfig(project_name="fake-project", dbt_project_path="/some/project/path")
execution_config = ExecutionConfig()
render_config = RenderConfig(enable_mock_profile=False)
profile_config = MagicMock()

with DAG("test-id", start_date=datetime(2024, 1, 1)) as dag:
DbtToAirflowConverter(
dag=dag,
nodes=nodes,
project_config=project_config,
profile_config=profile_config,
execution_config=execution_config,
render_config=render_config,
operator_args={},
)
task_args_cache_dir = mock_build_airflow_graph.call_args[1]["task_args"]["cache_dir"]
dbt_graph_cache_dir = mock_dbt_graph.call_args[1]["cache_dir"]

assert Path(tempfile.gettempdir()) in task_args_cache_dir.parents
assert task_args_cache_dir.parent.stem == "cosmos"
assert task_args_cache_dir.stem == "test-id"
assert task_args_cache_dir == dbt_graph_cache_dir


@patch("cosmos.settings.enable_cache", False)
@patch("cosmos.config.ProjectConfig.validate_project")
@patch("cosmos.converter.validate_initial_user_config")
@patch("cosmos.converter.DbtGraph")
@patch("cosmos.converter.build_airflow_graph")
def test_converter_disable_cache_sets_cache_dir_to_none(
mock_build_airflow_graph,
mock_dbt_graph,
mock_user_config,
mock_validate_project,
):
"""Tests that DbtGraph and operator and Airflow task args contain expected cache dir."""
project_config = ProjectConfig(project_name="fake-project", dbt_project_path="/some/project/path")
execution_config = ExecutionConfig()
render_config = RenderConfig(enable_mock_profile=False)
profile_config = MagicMock()

with DAG("test-id", start_date=datetime(2024, 1, 1)) as dag:
DbtToAirflowConverter(
dag=dag,
nodes=nodes,
project_config=project_config,
profile_config=profile_config,
execution_config=execution_config,
render_config=render_config,
operator_args={},
)
task_args_cache_dir = mock_build_airflow_graph.call_args[1]["task_args"]["cache_dir"]
dbt_graph_cache_dir = mock_dbt_graph.call_args[1]["cache_dir"]

assert dbt_graph_cache_dir is None
assert task_args_cache_dir == dbt_graph_cache_dir


@pytest.mark.parametrize(
"execution_mode,operator_args",
[
Expand Down

0 comments on commit e3fdce9

Please sign in to comment.