Repository: fal-ai/dbt-fal

Issue:
Better support for singular tests
<!-- *** Make sure you have searched for an existing bug report for this issue *** -->

**Describe the bug**
Fal object model does not seem to support singular tests. Getting various errors with dbt project that uses custom singular tests (standalone sql tests).

```py
from fal import FalDbt
faldbt = FalDbt(profiles_dir='.dbt', project_dir='/workspace')
print(faldbt.tests[0])
# Getting AttributeError: 'DbtTest' object has no attribute 'column'
for t in faldbt.tests:
    print(t.name)
# Geeting a AttributeError: 'ParsedSingularTestNode' object has no attribute 'test_metadata'
```

**Your environment**
```bash
# fal --version
fal 0.4.1
# dbt --version
Core:
  - installed: 1.2.0
  - latest:    1.2.0 - Up to date!
Plugins:
  - snowflake: 1.2.0 - Up to date!
```
**How to reproduce**
We have a dbt project with generic and custom tests, running `dbt run` and then code above

**Expected behavior**
Expect to be able to retrieve items from tests collection and be able to work with them regardless of type of test

**Actual behavior**
Can not work with tests in fal


Discussion:
Thanks for the report @miro-ur , I will take a look!

Resolving patch:
diff --git a/src/fal/__init__.py b/src/fal/__init__.py
--- a/src/fal/__init__.py
+++ b/src/fal/__init__.py
@@ -1,4 +1,11 @@
 from dbt.contracts.results import NodeStatus
 
-from faldbt.project import FalDbt, DbtModel, DbtSource, DbtTest
+from faldbt.project import (
+    FalDbt,
+    DbtModel,
+    DbtSource,
+    DbtTest,
+    DbtGenericTest,
+    DbtSingularTest,
+)
 from fal.fal_script import Context, CurrentModel
diff --git a/src/faldbt/project.py b/src/faldbt/project.py
--- a/src/faldbt/project.py
+++ b/src/faldbt/project.py
@@ -17,7 +17,12 @@
 from pathlib import Path
 
 from dbt.node_types import NodeType
-from dbt.contracts.graph.parsed import ParsedSourceDefinition, TestMetadata
+from dbt.contracts.graph.parsed import (
+    ParsedSourceDefinition,
+    TestMetadata,
+    ParsedGenericTestNode,
+    ParsedSingularTestNode,
+)
 from dbt.contracts.graph.compiled import ManifestNode
 from dbt.contracts.graph.manifest import (
     Manifest,
@@ -86,35 +91,64 @@ def _set_status(self, status: str):
 
 @dataclass
 class DbtTest(_DbtNode):
-    column: str = field(init=False)
-    source: Optional[Tuple[str, str]] = field(init=False, default=None)
-    model: Optional[str] = field(init=False, default=None)
+    model_ids: List[str] = field(init=False, default_factory=list)
+    source_ids: List[str] = field(init=False, default_factory=list)
 
-    def __post_init__(self):
-        node = self.node
+    @classmethod
+    def init(cls, node):
         if node.resource_type == NodeType.Test:
-            if hasattr(node, "test_metadata"):
-                metadata: TestMetadata = self.node.test_metadata
+            if isinstance(node, ParsedGenericTestNode):
+                test = DbtGenericTest(node=node)
+            elif isinstance(node, ParsedSingularTestNode):
+                test = DbtSingularTest(node=node)
+            else:
+                raise ValueError(f"Unexpected test class {node.__class__.__name__}")
 
-                # metadata.kwargs looks like this:
-                # kwargs={'column_name': 'y', 'model': "{{ get_where_subquery(ref('boston')) }}"}
-                # and we want to get 'boston' is the model name that we want extract
-                # except in dbt v 0.20.1, where we need to filter 'where' string out
-                self.column = metadata.kwargs.get("column_name", None)
+            for dep in test.node.depends_on.nodes:
+                if dep.startswith("model."):
+                    test.model_ids.append(dep)
+                if dep.startswith("source."):
+                    test.source_ids.append(dep)
 
-                # TODO: Handle package models?
-                refs = re.findall(r"ref\('([^']+)'\)", metadata.kwargs["model"])
+            return test
+        else:
+            raise TypeError(
+                f"Initialized DbtTest with node of type {node.resource_type}"
+            )
 
-                if refs:
-                    self.model = refs[0]
 
-                sources = re.findall(
-                    r"source\('([^']+)', *'([^']+)'\)", metadata.kwargs["model"]
-                )
-                if sources:
-                    self.source = sources[0]
-            else:
-                logger.warn(f"Non-generic test was not processed: {node.name}")
+@dataclass
+class DbtGenericTest(DbtTest):
+    column: Optional[str] = field(init=False)
+
+    def __post_init__(self):
+        assert isinstance(self.node, ParsedGenericTestNode)
+        self.column = self.node.column_name
+
+    @property
+    def source_id(self):
+        if self.source_ids:
+            return self.source_ids[0]
+
+    @property
+    def model_id(self):
+        if self.model_ids:
+            return self.model_ids[0]
+
+    # TODO: Deprecate?
+    @property
+    def source(self):
+        if self.source_id:
+            parts = self.source_id.split(".")
+            return parts[-2], parts[-1]
+
+    # TODO: Deprecate?
+    @property
+    def model(self):
+        if self.model_id:
+            # TODO: handle package models
+            parts = self.model_id.split(".")
+            return parts[-1]
 
     @property
     def name(self) -> str:
@@ -122,9 +156,16 @@ def name(self) -> str:
         return metadata.name
 
 
+@dataclass
+class DbtSingularTest(DbtTest):
+    def __post_init__(self):
+        assert isinstance(self.node, ParsedSingularTestNode)
+
+
 @dataclass
 class _DbtTestableNode(_DbtNode):
-    tests: List[DbtTest] = field(default_factory=list)
+    # TODO: should this include singular tests that ref to this node?
+    tests: List[DbtGenericTest] = field(default_factory=list)
 
     def _get_status(self):
         if self._status == NodeStatus.Skipped and any(
@@ -298,11 +339,10 @@ def _map_nodes(
         results_map = {r.unique_id: r for r in run_results.results}
 
         tests: List[DbtTest] = []
-        tests_dict: Dict[Union[Tuple[str, str], str], List[DbtTest]] = defaultdict(list)
+
+        tests_dict: Dict[str, List[DbtGenericTest]] = defaultdict(list)
         for node in self.get_test_nodes():
-            test = DbtTest(
-                node=node,
-            )
+            test: DbtTest = DbtTest.init(node=node)
 
             result = results_map.get(node.unique_id)
             if result:
@@ -310,16 +350,17 @@ def _map_nodes(
 
             tests.append(test)
 
-            if test.model:
-                tests_dict[test.model].append(test)
-            if test.source:
-                tests_dict[test.source].append(test)
+            if isinstance(test, DbtGenericTest):
+                if test.model_id:
+                    tests_dict[test.model_id].append(test)
+                if test.source_id:
+                    tests_dict[test.source_id].append(test)
 
         models: List[DbtModel] = []
         for node in self.get_model_nodes():
             model = DbtModel(
                 node=node,
-                tests=tests_dict[node.name],
+                tests=tests_dict[node.unique_id],
                 python_model=generated_models.get(node.name),
             )
 
@@ -336,7 +377,7 @@ def _map_nodes(
         for node in self.get_source_nodes():
             source = DbtSource(
                 node=node,
-                tests=tests_dict[(node.source_name, node.name)],
+                tests=tests_dict[node.unique_id],
                 freshness=source_freshness_map.get(node.unique_id),
             )
 


Test patch:
diff --git a/integration_tests/features/dbt_test.feature b/integration_tests/features/dbt_test.feature
--- a/integration_tests/features/dbt_test.feature
+++ b/integration_tests/features/dbt_test.feature
@@ -34,3 +34,24 @@ Feature: `dbt test` awareness
       | source results.other_source has 0 tests, source status is skipped |
       | model some_model has 2 tests, model status is tested              |
       | model other_model has 0 tests, model status is skipped            |
+
+  Scenario: Singular tests are present in context
+    Given the project 002_jaffle_shop
+    When the following shell command is invoked:
+      """
+      dbt seed --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    When the following command is invoked:
+      """
+      fal flow run --profiles-dir $profilesDir --project-dir $baseDir --exclude orders_forecast+
+      """
+    And the following shell command is invoked:
+      """
+      dbt test --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    And the following shell command is invoked:
+      """
+      python $baseDir/fal_dbt.py $baseDir $profilesDir
+      """
+    Then the script fal_dbt.py output file has the lines:
+      | There are 21 tests | There are 20 generic tests | There are 1 singular tests |
diff --git a/integration_tests/features/environment.py b/integration_tests/features/environment.py
--- a/integration_tests/features/environment.py
+++ b/integration_tests/features/environment.py
@@ -21,9 +21,11 @@ def after_scenario(context, scenario):
     if hasattr(context, "exc") and context.exc:
         from traceback import print_exception
 
-        print_exception(*context.exc)
         _etype, exception, _tb = context.exc
-        raise exception
+
+        print_exception(*context.exc)
+
+        raise AssertionError("Should have expected exception") from exception
 
 
 def before_tag(context, tag):
diff --git a/integration_tests/projects/002_jaffle_shop/fal_dbt.py b/integration_tests/projects/002_jaffle_shop/fal_dbt.py
new file mode 100644
--- /dev/null
+++ b/integration_tests/projects/002_jaffle_shop/fal_dbt.py
@@ -0,0 +1,36 @@
+import sys
+from fal import FalDbt
+from _fal_testing import create_file
+
+from faldbt.project import DbtGenericTest, DbtSingularTest
+
+project_dir = "."
+if len(sys.argv) >= 2:
+    project_dir = sys.argv[1]
+profiles_dir = "~/.dbt"
+if len(sys.argv) >= 3:
+    profiles_dir = sys.argv[2]
+
+print(f"project_dir={project_dir}, profiles_dir={profiles_dir}")
+
+faldbt = FalDbt(project_dir=project_dir, profiles_dir=profiles_dir)
+print(faldbt)
+print(faldbt.tests)
+
+for test in faldbt.tests:
+    print(test)
+
+for source in faldbt.sources:
+    print(source.name, source.table_name, [(t.name, t.status) for t in source.tests])
+
+for model in faldbt.models:
+    print(model.name, [(t.name, t.status) for t in model.tests])
+
+output = f"There are {len(faldbt.tests)} tests\n"
+output += f"There are {len([t for t in faldbt.tests if isinstance(t, DbtGenericTest)])} generic tests\n"
+output += f"There are {len([t for t in faldbt.tests if isinstance(t, DbtSingularTest)])} singular tests\n"
+
+for test in faldbt.tests:
+    output += f"test {test.name} {'generic' if isinstance(test, DbtGenericTest) else 'singular'} \n"
+
+create_file(output, "fal_dbt.txt")
diff --git a/integration_tests/projects/002_jaffle_shop/tests/customer_lifetime_value.sql b/integration_tests/projects/002_jaffle_shop/tests/customer_lifetime_value.sql
new file mode 100644
--- /dev/null
+++ b/integration_tests/projects/002_jaffle_shop/tests/customer_lifetime_value.sql
@@ -0,0 +1,4 @@
+SELECT
+    customer_id
+FROM {{ ref('customers') }}
+WHERE NOT (customer_lifetime_value >= 0) OR customer_lifetime_value IS NULL
diff --git a/tests/_fal_testing/utils.py b/tests/_fal_testing/utils.py
--- a/tests/_fal_testing/utils.py
+++ b/tests/_fal_testing/utils.py
@@ -8,14 +8,18 @@ def create_artifact(context, suffix, additional_data=None):
     model_name = model.name if model else "GLOBAL"
     model_status = model.status if model else None
 
-    temp_dir = Path(os.environ["temp_dir"])
-    temp_file = (temp_dir / model_name).with_suffix(suffix)
-
     output = f"Model name: {model_name}"
     output += f"\nStatus: {model_status}"
     output += f"\nModel dataframe name: {model_name}"
     if additional_data:
         output += f"\n{additional_data}"
+
+    create_file(output, Path(model_name).with_suffix(suffix))
+
+
+def create_file(output, file_name):
+    temp_dir = Path(os.environ["temp_dir"])
+    temp_file = temp_dir / file_name
     temp_file.write_text(output)