From ec7e502f1b6ceb045cdb7e78fd2a85ebd478df41 Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Wed, 24 Apr 2024 10:53:07 +0100
Subject: [PATCH] Add test to confirm behaviour of converter on cache dir

---
 cosmos/settings.py      |  1 +
 tests/test_converter.py | 70 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/cosmos/settings.py b/cosmos/settings.py
index c889b04b8..9670ee7d1 100644
--- a/cosmos/settings.py
+++ b/cosmos/settings.py
@@ -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
diff --git a/tests/test_converter.py b/tests/test_converter.py
index f54a95fed..bef2dc06d 100644
--- a/tests/test_converter.py
+++ b/tests/test_converter.py
@@ -1,3 +1,4 @@
+import tempfile
 from datetime import datetime
 from pathlib import Path
 from unittest.mock import MagicMock, patch
@@ -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",
     [