From 5c010cc430f1e81d50063c36195fa655e1393ff5 Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Thu, 18 Apr 2024 23:37:54 +0100
Subject: [PATCH] Increase test coverage and fix implementation

---
 cosmos/cache.py                     |  10 +-
 cosmos/config.py                    |   3 +-
 cosmos/converter.py                 |  12 --
 cosmos/dbt/graph.py                 |   4 +-
 cosmos/operators/local.py           |   2 +-
 dev/dags/basic_cosmos_task_group.py |  10 +-
 tests/dbt/test_graph.py             |  18 ++-
 tests/plugin/test_plugin.py         | 221 ----------------------------
 tests/test_cache.py                 |  62 ++++++++
 9 files changed, 97 insertions(+), 245 deletions(-)
 delete mode 100644 tests/plugin/test_plugin.py
 create mode 100644 tests/test_cache.py

diff --git a/cosmos/cache.py b/cosmos/cache.py
index 01db22c0f..7d24b5ab0 100644
--- a/cosmos/cache.py
+++ b/cosmos/cache.py
@@ -11,6 +11,11 @@
 from cosmos.dbt.project import get_partial_parse_path
 
 
+# It was considered to create a cache identifier based on the dbt project path, as opposed
+# to where it is used in Airflow. However, we could have concurrency issues if the same
+# dbt cached directory was being used by different dbt task groups or DAGs within the same
+# node. For this reason, as a starting point, the cache is identified by where it is used.
+# This can be reviewed in the future.
 def create_cache_identifier(dag: DAG, task_group: TaskGroup | None) -> str:
     """
     Given a DAG name and a (optional) task_group_name, create the identifier for caching.
@@ -22,11 +27,8 @@ def create_cache_identifier(dag: DAG, task_group: TaskGroup | None) -> str:
     if task_group:
         if task_group.dag_id is not None:
             cache_identifiers_list = [task_group.dag_id]
-        if task_group.upstream_group_ids is not None:
-            group_ids: list[str] = [tg for tg in task_group.upstream_group_ids or [] if tg is not None]
-            cache_identifiers_list.extend(group_ids)
         if task_group.group_id is not None:
-            cache_identifiers_list.extend(task_group.group_id)
+            cache_identifiers_list.extend([task_group.group_id.replace(".", "_")])
         cache_identifier = "_".join(cache_identifiers_list)
     else:
         cache_identifier = dag.dag_id
diff --git a/cosmos/config.py b/cosmos/config.py
index 18311057c..b5346377b 100644
--- a/cosmos/config.py
+++ b/cosmos/config.py
@@ -288,7 +288,8 @@ def ensure_profile(
                 with tempfile.TemporaryDirectory() as temp_dir:
                     temp_file = Path(temp_dir) / DEFAULT_PROFILES_FILE_NAME
                     logger.info(
-                        "Creating temporary profiles.yml at %s with the following contents:\n%s",
+                        "Creating temporary profiles.yml with use_mock_values=%s at %s with the following contents:\n%s",
+                        use_mock_values,
                         temp_file,
                         profile_contents,
                     )
diff --git a/cosmos/converter.py b/cosmos/converter.py
index 273519be4..9b33efa20 100644
--- a/cosmos/converter.py
+++ b/cosmos/converter.py
@@ -215,8 +215,6 @@ def __init__(
 
         validate_initial_user_config(execution_config, profile_config, project_config, render_config, operator_args)
 
-        # If we are using the old interface, we should migrate it to the new interface
-        # This is safe to do now since we have validated which config interface we're using
         if project_config.dbt_project_path:
             execution_config, render_config = migrate_to_new_interface(execution_config, project_config, render_config)
 
@@ -227,16 +225,6 @@ def __init__(
 
         cache_dir = cache.obtain_cache_dir_path(cache_identifier=cache.create_cache_identifier(dag, task_group))
 
-        # Previously, we were creating a cosmos.dbt.project.DbtProject
-        # DbtProject has now been replaced with ProjectConfig directly
-        #   since the interface of the two classes were effectively the same
-        # Under this previous implementation, we were passing:
-        #  - name, root dir, models dir, snapshots dir and manifest path
-        # Internally in the dbtProject class, we were defaulting the profile_path
-        #   To be root dir/profiles.yml
-        # To keep this logic working, if converter is given no ProfileConfig,
-        #   we can create a default retaining this value to preserve this functionality.
-        # We may want to consider defaulting this value in our actual ProjectConfig class?
         self.dbt_graph = DbtGraph(
             project=project_config,
             render_config=render_config,
diff --git a/cosmos/dbt/graph.py b/cosmos/dbt/graph.py
index 597150ffe..80276af9f 100644
--- a/cosmos/dbt/graph.py
+++ b/cosmos/dbt/graph.py
@@ -74,7 +74,7 @@ def name(self) -> str:
 def run_command(command: list[str], tmp_dir: Path, env_vars: dict[str, str]) -> str:
     """Run a command in a subprocess, returning the stdout."""
     logger.info("Running command: `%s`", " ".join(command))
-    logger.info("Environment variable keys: %s", env_vars.keys())
+    logger.debug("Environment variable keys: %s", env_vars.keys())
     process = Popen(
         command,
         stdout=PIPE,
@@ -215,7 +215,7 @@ def run_dbt_ls(
 
         stdout = run_command(ls_command, tmp_dir, env_vars)
 
-        logger.debug("dbt ls output: %s", stdout)
+        logger.info("dbt ls output: %s", stdout)
         log_filepath = self.log_dir / DBT_LOG_FILENAME
         logger.debug("dbt logs available in: %s", log_filepath)
         if log_filepath.exists():
diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
index e28045ba2..5079fbaf4 100644
--- a/cosmos/operators/local.py
+++ b/cosmos/operators/local.py
@@ -325,7 +325,7 @@ def run_command(
 
                 full_cmd = cmd + flags
 
-                logger.info("Using environment variables keys: %s", env.keys())
+                logger.debug("Using environment variables keys: %s", env.keys())
 
                 result = self.invoke_dbt(
                     command=full_cmd,
diff --git a/dev/dags/basic_cosmos_task_group.py b/dev/dags/basic_cosmos_task_group.py
index 14d209d13..a06cb51be 100644
--- a/dev/dags/basic_cosmos_task_group.py
+++ b/dev/dags/basic_cosmos_task_group.py
@@ -25,7 +25,10 @@
     ),
 )
 
-shared_execution_config = ExecutionConfig(invocation_mode=InvocationMode.DBT_RUNNER)
+shared_execution_config = ExecutionConfig(
+    invocation_mode=InvocationMode.SUBPROCESS,
+    # invocation_mode=InvocationMode.DBT_RUNNER
+)
 
 
 @dag(
@@ -56,7 +59,10 @@ def basic_cosmos_task_group() -> None:
         project_config=ProjectConfig(
             (DBT_ROOT_PATH / "jaffle_shop").as_posix(),
         ),
-        render_config=RenderConfig(select=["path:seeds/raw_orders.csv"]),
+        render_config=RenderConfig(
+            select=["path:seeds/raw_orders.csv"],
+            enable_mock_profile=False,  # This is necessary to benefit from partial parsing when using ProfileMapping
+        ),
         execution_config=shared_execution_config,
         operator_args={"install_deps": True},
         profile_config=profile_config,
diff --git a/tests/dbt/test_graph.py b/tests/dbt/test_graph.py
index f72bbb146..faddf4484 100644
--- a/tests/dbt/test_graph.py
+++ b/tests/dbt/test_graph.py
@@ -482,7 +482,9 @@ def test_load_via_dbt_ls_without_dbt_deps(postgres_profile_config):
 
 
 @pytest.mark.integration
-def test_load_via_dbt_ls_without_dbt_deps_and_preinstalled_dbt_packages(tmp_dbt_project_dir, postgres_profile_config):
+def test_load_via_dbt_ls_without_dbt_deps_and_preinstalled_dbt_packages(
+    tmp_dbt_project_dir, postgres_profile_config, caplog, tmp_path
+):
     local_flags = [
         "--project-dir",
         tmp_dbt_project_dir / DBT_PROJECT_NAME,
@@ -506,17 +508,29 @@ def test_load_via_dbt_ls_without_dbt_deps_and_preinstalled_dbt_packages(tmp_dbt_
     stdout, stderr = process.communicate()
 
     project_config = ProjectConfig(dbt_project_path=tmp_dbt_project_dir / DBT_PROJECT_NAME)
-    render_config = RenderConfig(dbt_project_path=tmp_dbt_project_dir / DBT_PROJECT_NAME, dbt_deps=False)
+    render_config = RenderConfig(
+        dbt_project_path=tmp_dbt_project_dir / DBT_PROJECT_NAME, dbt_deps=False, enable_mock_profile=False
+    )
     execution_config = ExecutionConfig(dbt_project_path=tmp_dbt_project_dir / DBT_PROJECT_NAME)
     dbt_graph = DbtGraph(
         project=project_config,
         render_config=render_config,
         execution_config=execution_config,
         profile_config=postgres_profile_config,
+        cache_dir=tmp_path,
     )
 
+    from cosmos.constants import DBT_TARGET_DIR_NAME
+
+    (tmp_path / DBT_TARGET_DIR_NAME).mkdir(parents=True, exist_ok=True)
+
     dbt_graph.load_via_dbt_ls()  # does not raise exception
 
+    assert "Unable to do partial parsing" in caplog.text
+    # TODO: split the caching test into a separate test, and make the following assertion work
+    # dbt_graph.load_via_dbt_ls()  # does not raise exception
+    # assert not "Unable to do partial parsing" in caplog.text
+
 
 @pytest.mark.integration
 @patch("cosmos.dbt.graph.Popen")
diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py
deleted file mode 100644
index 8d0e3742f..000000000
--- a/tests/plugin/test_plugin.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# dbt-core relies on Jinja2>3, whereas Flask<2 relies on an incompatible version of Jinja2.
-#
-# This discrepancy causes the automated integration tests to fail, as dbt-core is installed in the same
-# environment as apache-airflow.
-#
-# We can get around this by patching the jinja2 namespace to include the deprecated objects:
-try:
-    import flask  # noqa: F401
-except ImportError:
-    import jinja2
-    import markupsafe
-
-    jinja2.Markup = markupsafe.Markup
-    jinja2.escape = markupsafe.escape
-
-import sys
-from unittest.mock import MagicMock, PropertyMock, mock_open, patch
-
-import pytest
-from airflow.configuration import conf
-from airflow.exceptions import AirflowConfigException
-from airflow.utils.db import initdb, resetdb
-from airflow.www.app import cached_app
-from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
-from flask.testing import FlaskClient
-
-import cosmos.plugin
-from cosmos.plugin import (
-    dbt_docs_view,
-    iframe_script,
-    open_azure_file,
-    open_file,
-    open_gcs_file,
-    open_http_file,
-    open_s3_file,
-)
-
-original_conf_get = conf.get
-
-
-def _get_text_from_response(response) -> str:
-    # Airflow < 2.4 uses an old version of Werkzeug that does not have Response.text.
-    if not hasattr(response, "text"):
-        return response.get_data(as_text=True)
-    else:
-        return response.text
-
-
-@pytest.fixture(scope="module")
-def app() -> FlaskClient:
-    initdb()
-
-    app = cached_app(testing=True)
-    appbuilder: AirflowAppBuilder = app.extensions["appbuilder"]
-
-    appbuilder.sm.check_authorization = lambda *args, **kwargs: True
-
-    if dbt_docs_view not in appbuilder.baseviews:
-        appbuilder._check_and_init(dbt_docs_view)
-        appbuilder.register_blueprint(dbt_docs_view)
-
-    yield app.test_client()
-
-    resetdb(skip_init=True)
-
-
-def test_dbt_docs(monkeypatch, app):
-    def conf_get(section, key, *args, **kwargs):
-        if section == "cosmos" and key == "dbt_docs_dir":
-            return "path/to/docs/dir"
-        else:
-            return original_conf_get(section, key, *args, **kwargs)
-
-    monkeypatch.setattr(conf, "get", conf_get)
-
-    response = app.get("/cosmos/dbt_docs")
-
-    assert response.status_code == 200
-    assert "<iframe" in _get_text_from_response(response)
-
-
-@patch.object(cosmos.plugin, "open_file")
-def test_dbt_docs_not_set_up(monkeypatch, app):
-    response = app.get("/cosmos/dbt_docs")
-
-    assert response.status_code == 200
-    assert "<iframe" not in _get_text_from_response(response)
-
-
-@patch.object(cosmos.plugin, "open_file")
-@pytest.mark.parametrize("artifact", ["dbt_docs_index.html", "manifest.json", "catalog.json"])
-def test_dbt_docs_artifact(mock_open_file, monkeypatch, app, artifact):
-    def conf_get(section, key, *args, **kwargs):
-        if section == "cosmos" and key == "dbt_docs_dir":
-            return "path/to/docs/dir"
-        else:
-            return original_conf_get(section, key, *args, **kwargs)
-
-    monkeypatch.setattr(conf, "get", conf_get)
-
-    if artifact == "dbt_docs_index.html":
-        mock_open_file.return_value = "<head></head><body></body>"
-    else:
-        mock_open_file.return_value = "{}"
-
-    response = app.get(f"/cosmos/{artifact}")
-
-    mock_open_file.assert_called_once()
-    assert response.status_code == 200
-    if artifact == "dbt_docs_index.html":
-        assert iframe_script in _get_text_from_response(response)
-
-
-@pytest.mark.parametrize("artifact", ["dbt_docs_index.html", "manifest.json", "catalog.json"])
-def test_dbt_docs_artifact_missing(app, artifact):
-    response = app.get(f"/cosmos/{artifact}")
-    assert response.status_code == 404
-
-
-@pytest.mark.parametrize(
-    "path,open_file_callback",
-    [
-        ("s3://my-bucket/my/path/", "open_s3_file"),
-        ("gs://my-bucket/my/path/", "open_gcs_file"),
-        ("wasb://my-bucket/my/path/", "open_azure_file"),
-        ("http://my-bucket/my/path/", "open_http_file"),
-        ("https://my-bucket/my/path/", "open_http_file"),
-    ],
-)
-def test_open_file_calls(path, open_file_callback, monkeypatch):
-    def conf_get(section, key, *args, **kwargs):
-        if section == "cosmos" and key == "dbt_docs_conn_id":
-            return "mock_conn_id"
-        else:
-            return original_conf_get(section, key, *args, **kwargs)
-
-    monkeypatch.setattr(conf, "get", conf_get)
-
-    with patch.object(cosmos.plugin, open_file_callback) as mock_callback:
-        mock_callback.return_value = "mock file contents"
-        res = open_file(path)
-
-    mock_callback.assert_called_with(conn_id="mock_conn_id", path=path)
-    assert res == "mock file contents"
-
-
-@pytest.mark.parametrize("conn_id", ["mock_conn_id", None])
-def test_open_s3_file(conn_id):
-    mock_module = MagicMock()
-    with patch.dict(sys.modules, {"airflow.providers.amazon.aws.hooks.s3": mock_module}):
-        mock_hook = mock_module.S3Hook.return_value
-        mock_hook.read_key.return_value = "mock file contents"
-
-        res = open_s3_file(conn_id=conn_id, path="s3://mock-path/to/docs")
-
-        if conn_id is not None:
-            mock_module.S3Hook.assert_called_once_with(aws_conn_id=conn_id)
-
-        mock_hook.read_key.assert_called_once_with(bucket_name="mock-path", key="to/docs")
-        assert res == "mock file contents"
-
-
-@pytest.mark.parametrize("conn_id", ["mock_conn_id", None])
-def test_open_gcs_file(conn_id):
-    mock_module = MagicMock()
-    with patch.dict(sys.modules, {"airflow.providers.google.cloud.hooks.gcs": mock_module}):
-        mock_hook = mock_module.GCSHook.return_value = MagicMock()
-        mock_hook.download.return_value = b"mock file contents"
-
-        res = open_gcs_file(conn_id=conn_id, path="gs://mock-path/to/docs")
-
-        if conn_id is not None:
-            mock_module.GCSHook.assert_called_once_with(gcp_conn_id=conn_id)
-
-        mock_hook.download.assert_called_once_with(bucket_name="mock-path", object_name="to/docs")
-        assert res == "mock file contents"
-
-
-@pytest.mark.parametrize("conn_id", ["mock_conn_id", None])
-def test_open_azure_file(conn_id):
-    mock_module = MagicMock()
-    with patch.dict(sys.modules, {"airflow.providers.microsoft.azure.hooks.wasb": mock_module}):
-        mock_hook = mock_module.WasbHook.return_value = MagicMock()
-        mock_hook.default_conn_name = PropertyMock(return_value="default_conn")
-        mock_hook.read_file.return_value = "mock file contents"
-
-        res = open_azure_file(conn_id=conn_id, path="wasb://mock-path/to/docs")
-
-        if conn_id is not None:
-            mock_module.WasbHook.assert_called_once_with(wasb_conn_id=conn_id)
-
-        mock_hook.read_file.assert_called_once_with(container_name="mock-path", blob_name="to/docs")
-        assert res == "mock file contents"
-
-
-@pytest.mark.parametrize("conn_id", ["mock_conn_id", None])
-def test_open_http_file(conn_id):
-    mock_module = MagicMock()
-    with patch.dict(sys.modules, {"airflow.providers.http.hooks.http": mock_module}):
-        mock_hook = mock_module.HttpHook.return_value = MagicMock()
-        mock_response = mock_hook.run.return_value = MagicMock()
-        mock_hook.default_conn_name = PropertyMock(return_value="default_conn")
-        mock_hook.check_response.return_value = mock_response
-        mock_response.text = "mock file contents"
-
-        res = open_http_file(conn_id=conn_id, path="http://mock-path/to/docs")
-
-        if conn_id is not None:
-            mock_module.HttpHook.assert_called_once_with(method="GET", http_conn_id=conn_id)
-        else:
-            mock_module.HttpHook.assert_called_once_with(method="GET", http_conn_id="")
-
-        mock_hook.run.assert_called_once_with(endpoint="http://mock-path/to/docs")
-        assert res == "mock file contents"
-
-
-@patch("builtins.open", new_callable=mock_open, read_data="mock file contents")
-def test_open_file_local(mock_file):
-    res = open_file("/my/path")
-    mock_file.assert_called_with("/my/path")
-    assert res == "mock file contents"
diff --git a/tests/test_cache.py b/tests/test_cache.py
new file mode 100644
index 000000000..b8314e19d
--- /dev/null
+++ b/tests/test_cache.py
@@ -0,0 +1,62 @@
+import os
+from datetime import datetime
+
+import pytest
+from airflow import DAG
+from airflow.utils.task_group import TaskGroup
+
+from cosmos.cache import create_cache_identifier, get_latest_partial_parse
+from cosmos.constants import DBT_PARTIAL_PARSE_FILE_NAME, DBT_TARGET_DIR_NAME
+
+START_DATE = datetime(2024, 4, 16)
+example_dag = DAG("dag", start_date=START_DATE)
+
+
+@pytest.mark.parametrize(
+    "dag, task_group, result_identifier",
+    [
+        (example_dag, None, "dag"),
+        (None, TaskGroup(dag=example_dag, group_id="inner_tg"), "dag_inner_tg"),
+        (
+            None,
+            TaskGroup(
+                dag=example_dag, group_id="child_tg", parent_group=TaskGroup(dag=example_dag, group_id="parent_tg")
+            ),
+            "dag_parent_tg_child_tg",
+        ),
+        (
+            None,
+            TaskGroup(
+                dag=example_dag,
+                group_id="child_tg",
+                parent_group=TaskGroup(
+                    dag=example_dag, group_id="mum_tg", parent_group=TaskGroup(dag=example_dag, group_id="nana_tg")
+                ),
+            ),
+            "dag_nana_tg_mum_tg_child_tg",
+        ),
+    ],
+)
+def test_create_cache_identifier(dag, task_group, result_identifier):
+    assert create_cache_identifier(dag, task_group) == result_identifier
+
+
+def test_get_latest_partial_parse(tmp_path):
+    old_tmp_dir = tmp_path / "old"
+    old_tmp_target_dir = old_tmp_dir / DBT_TARGET_DIR_NAME
+    old_tmp_target_dir.mkdir(create_parents=True, exist_ok=True)
+    old_partial_parse_filepath = old_tmp_target_dir / DBT_PARTIAL_PARSE_FILE_NAME
+    open(old_partial_parse_filepath, "a").close()
+
+    new_tmp_dir = tmp_path / "new"
+    new_tmp_target_dir = new_tmp_dir / DBT_TARGET_DIR_NAME
+    new_tmp_target_dir.mkdir(create_parents=True, exist_ok=True)
+    new_partial_parse_filepath = new_tmp_target_dir / DBT_PARTIAL_PARSE_FILE_NAME
+    open(new_partial_parse_filepath, "a").close()
+
+    assert get_latest_partial_parse(old_tmp_dir, new_tmp_dir) == new_partial_parse_filepath
+    assert get_latest_partial_parse(new_tmp_dir, old_tmp_dir) == new_partial_parse_filepath
+    assert get_latest_partial_parse(old_tmp_dir, old_tmp_dir) == old_partial_parse_filepath
+    assert get_latest_partial_parse(old_tmp_dir, tmp_path) == old_partial_parse_filepath
+    assert get_latest_partial_parse(tmp_path, old_tmp_dir) == old_partial_parse_filepath
+    assert get_latest_partial_parse(tmp_path, tmp_path) is None