From 684a8caf2eb3fa9837909d9d03e465bbd36317ab Mon Sep 17 00:00:00 2001
From: Dave Connors <dave.connors@fishtownanalytics.com>
Date: Wed, 29 May 2024 08:26:02 -0500
Subject: [PATCH] add custom generic test and jinja handling logic

---
 dbt_meshify/dbt_projects.py                          |  9 +++++++--
 .../split/split_proj/models/marts/__models.yml       |  1 +
 .../split_proj/tests/generic/custom_generic_test.sql |  3 +++
 tests/unit/test_jinja_blocks.py                      | 12 ++++++++++++
 4 files changed, 23 insertions(+), 2 deletions(-)
 create mode 100644 test-projects/split/split_proj/tests/generic/custom_generic_test.sql

diff --git a/dbt_meshify/dbt_projects.py b/dbt_meshify/dbt_projects.py
index fe69bcf..f2e3c2f 100644
--- a/dbt_meshify/dbt_projects.py
+++ b/dbt_meshify/dbt_projects.py
@@ -140,7 +140,7 @@ def models(self) -> Dict[str, ModelNode]:
             return self._models
 
         self._models = {
-            node_name: node
+            node_name: node  # type: ignore
             for node_name, node in self.manifest.nodes.items()
             if node.resource_type == "model" and isinstance(node, ModelNode)
         }
@@ -328,11 +328,16 @@ def find_jinja_blocks(self) -> Dict[str, JinjaBlock]:
             )
 
         for unique_id, macro in self.manifest.macros.items():
+            block_type = "macro"
+            name = macro.name
             if macro.package_name != self.name:
                 continue
+            if "tests/generic/" in macro.path:
+                block_type = "test"
+                name = macro.name[5:]
 
             blocks[unique_id] = JinjaBlock.from_file(
-                path=self.path / macro.original_file_path, block_type="macro", name=macro.name
+                path=self.path / macro.original_file_path, block_type=block_type, name=name
             )
 
         return blocks
diff --git a/test-projects/split/split_proj/models/marts/__models.yml b/test-projects/split/split_proj/models/marts/__models.yml
index 685b402..e095108 100644
--- a/test-projects/split/split_proj/models/marts/__models.yml
+++ b/test-projects/split/split_proj/models/marts/__models.yml
@@ -45,6 +45,7 @@ models:
           expression: count_food_items + count_drink_items = count_items
       - dbt_utils.expression_is_true:
           expression: subtotal_food_items + subtotal_drink_items = subtotal
+      - custom_generic_test
     columns:
       - name: order_id
         description: The unique key of the orders mart.
diff --git a/test-projects/split/split_proj/tests/generic/custom_generic_test.sql b/test-projects/split/split_proj/tests/generic/custom_generic_test.sql
new file mode 100644
index 0000000..0269fe1
--- /dev/null
+++ b/test-projects/split/split_proj/tests/generic/custom_generic_test.sql
@@ -0,0 +1,3 @@
+{% test custom_generic_test() %}
+  select true where false
+{% endtest %}
\ No newline at end of file
diff --git a/tests/unit/test_jinja_blocks.py b/tests/unit/test_jinja_blocks.py
index 9c62367..1ad60e9 100644
--- a/tests/unit/test_jinja_blocks.py
+++ b/tests/unit/test_jinja_blocks.py
@@ -68,6 +68,14 @@
 {% endmacro %}
 """
 
+simple_custom_generic_test = """\
+
+
+{% test my_custom_test(model) %}
+  select * from {{ model }} where false
+{% endtest %}
+"""
+
 simple_macro_no_spaces = """\
 
 
@@ -176,3 +184,7 @@ def test_from_file_extracts_content_in_files_with_multiple_blocks(self):
             content
             == "{% docs potato_name %}\nThe name of the customer's favorite potato dish.\n{% enddocs %}"
         )
+
+    def test_from_file_extracts_custom_generic_test(self):
+        range = JinjaBlock.find_block_range(simple_custom_generic_test, "test", "my_custom_test")
+        assert range == (2, 88)