From 58de67e34243ea38024461173034e9b30c6ee700 Mon Sep 17 00:00:00 2001
From: edgga <65926599+edgga@users.noreply.github.com>
Date: Mon, 30 Oct 2023 16:35:44 +0200
Subject: [PATCH] Fix adding test nodes to DAGs built using
 `LoadMethod.DBT_MANIFEST` and `LoadMethod.CUSTOM` (#615)

Resolving issues with the DBT_MANIFEST/CUSTOM load methods when the
has_test attribute is not assigned to the node correctly.

## Description

When a tag selector is used, all tests are filtered out because of the
DbtResourceType.TEST node does not have any information about tags. To
bypass this limitation - tags are assigned to tests based on their
parent model.

## Related Issue(s)

Closes: #580

Co-authored-by: edgarasnavickas <edgaras.navickas@kevin.eu>
---
 cosmos/dbt/selector.py     |  3 +++
 tests/dbt/test_graph.py    | 20 ++++++++++++++++++++
 tests/sample/manifest.json |  8 ++++++--
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/cosmos/dbt/selector.py b/cosmos/dbt/selector.py
index 0e6034a0b..f195f225b 100644
--- a/cosmos/dbt/selector.py
+++ b/cosmos/dbt/selector.py
@@ -108,6 +108,9 @@ def should_include_node(node_id: str, node: DbtNode) -> bool:
 
         visited_nodes.add(node_id)
 
+        if node.resource_type == DbtResourceType.TEST:
+            node.tags = getattr(nodes.get(node.depends_on[0]), "tags", [])
+
         if config.tags:
             if not (set(config.tags) <= set(node.tags)):
                 return False
diff --git a/tests/dbt/test_graph.py b/tests/dbt/test_graph.py
index 3927bbfdd..b36bf92fb 100644
--- a/tests/dbt/test_graph.py
+++ b/tests/dbt/test_graph.py
@@ -504,6 +504,26 @@ def test_update_node_dependency_test_not_exist():
         assert nodes.has_test is False
 
 
+def test_tag_selected_node_test_exist():
+    project_config = ProjectConfig(
+        dbt_project_path=DBT_PROJECTS_ROOT_DIR / DBT_PROJECT_NAME, manifest_path=SAMPLE_MANIFEST
+    )
+    profile_config = ProfileConfig(
+        profile_name="test",
+        target_name="test",
+        profiles_yml_filepath=DBT_PROJECTS_ROOT_DIR / DBT_PROJECT_NAME / "profiles.yml",
+    )
+    dbt_graph = DbtGraph(project=project_config, profile_config=profile_config, select=["tag:test_tag"])
+    dbt_graph.load_from_dbt_manifest()
+
+    assert len(dbt_graph.filtered_nodes) > 0
+
+    for _, node in dbt_graph.filtered_nodes.items():
+        assert node.tags == ["test_tag"]
+        if node.resource_type == DbtResourceType.MODEL:
+            assert node.has_test is True
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("load_method", ["load_via_dbt_ls", "load_from_dbt_manifest"])
 def test_load_dbt_ls_and_manifest_with_model_version(load_method):
diff --git a/tests/sample/manifest.json b/tests/sample/manifest.json
index 4150234b8..d0b19c7b6 100644
--- a/tests/sample/manifest.json
+++ b/tests/sample/manifest.json
@@ -7576,7 +7576,9 @@
       "resource_type": "model",
       "schema": "public",
       "sources": [],
-      "tags": [],
+      "tags": [
+        "test_tag"
+      ],
       "unique_id": "model.jaffle_shop.customers",
       "unrendered_config": {
         "materialized": "table"
@@ -7754,7 +7756,9 @@
       "resource_type": "model",
       "schema": "public",
       "sources": [],
-      "tags": [],
+      "tags": [
+        "test_tag"
+      ],
       "unique_id": "model.jaffle_shop.orders",
       "unrendered_config": {
         "materialized": "table"