Repository: TobikoData/sqlmesh
Pull request #2271 (2024-03-14 21:23:29)

Feat: Record successful tests and return them in the API


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 4f053650..aab35699 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -94,6 +94,7 @@ from sqlmesh.core.state_sync import (
 )
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
+    ModelTextTestResult,
     generate_test,
     get_all_model_tests,
     run_model_tests,
@@ -1307,7 +1308,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         tests: t.Optional[t.List[str]] = None,
         verbose: bool = False,
         stream: t.Optional[t.TextIO] = None,
-    ) -> unittest.result.TestResult:
+    ) -> ModelTextTestResult:
         """Discover and run model tests"""
         if verbose:
             pd.set_option("display.max_columns", None)
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 30f65bc4..11ae7598 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -425,12 +425,17 @@ class _Model(ModelMeta, frozen=True):
         query = self.render_query_or_raise(**render_kwarg).copy()
 
         for select_or_union in query.find_all(exp.Select, exp.Union):
+            cte = select_or_union.find_ancestor(exp.With, exp.Select, exp.Subquery)
+            skip_limit = isinstance(cte, exp.With) and cte.recursive
+
             if isinstance(select_or_union, exp.Select) and select_or_union.args.get("from"):
                 select_or_union.where(exp.false(), copy=False)
-                if not isinstance(select_or_union.parent, exp.Union):
+                if not skip_limit and not isinstance(select_or_union.parent, exp.Union):
                     select_or_union.limit(0, copy=False)
-            elif isinstance(select_or_union, exp.Union) and not isinstance(
-                select_or_union.parent, exp.Union
+            elif (
+                not skip_limit
+                and isinstance(select_or_union, exp.Union)
+                and not isinstance(select_or_union.parent, exp.Union)
             ):
                 select_or_union.set("limit", exp.Limit(expression=exp.Literal.number(0)))
 
diff --git a/web/server/api/endpoints/commands.py b/web/server/api/endpoints/commands.py
index 7f24e584..5e89d47a 100644
--- a/web/server/api/endpoints/commands.py
+++ b/web/server/api/endpoints/commands.py
@@ -183,6 +183,13 @@ async def test(
                 (t.cast(ModelTest, test), reason) for test, reason in result.skipped
             )
         ],
+        successes=[
+            models.TestCase(
+                name=test.test_name,
+                path=test.path_relative_to(context.path),
+            )
+            for test in (t.cast(ModelTest, test) for test in result.successes)
+        ],
         tests_run=result.testsRun,
     )
 
diff --git a/web/server/models.py b/web/server/models.py
index 0eb90b82..2198a599 100644
--- a/web/server/models.py
+++ b/web/server/models.py
@@ -432,6 +432,7 @@ class TestResult(BaseModel):
     failures: t.List[TestErrorOrFailure]
     errors: t.List[TestErrorOrFailure]
     skipped: t.List[TestSkipped]
+    successes: t.List[TestCase]
 
 
 class ArtifactChange(BaseModel):


Test patch:
diff --git a/sqlmesh/core/test/__init__.py b/sqlmesh/core/test/__init__.py
index d5d44ca7..cc1df37f 100644
--- a/sqlmesh/core/test/__init__.py
+++ b/sqlmesh/core/test/__init__.py
@@ -25,7 +25,7 @@ def run_tests(
     verbosity: int = 1,
     stream: t.TextIO | None = None,
     default_catalog: str | None = None,
-) -> unittest.result.TestResult:
+) -> ModelTextTestResult:
     """Create a test suite of ModelTest objects and run it.
 
     Args:
@@ -47,9 +47,12 @@ def run_tests(
         for metadata in model_test_metadata
     )
 
-    return unittest.TextTestRunner(
-        stream=stream, verbosity=verbosity, resultclass=ModelTextTestResult
-    ).run(suite)
+    return t.cast(
+        ModelTextTestResult,
+        unittest.TextTestRunner(
+            stream=stream, verbosity=verbosity, resultclass=ModelTextTestResult
+        ).run(suite),
+    )
 
 
 def run_model_tests(
@@ -61,7 +64,7 @@ def run_model_tests(
     patterns: list[str] | None = None,
     stream: t.TextIO | None = None,
     default_catalog: t.Optional[str] = None,
-) -> unittest.result.TestResult:
+) -> ModelTextTestResult:
     """Load and run tests.
 
     Args
diff --git a/sqlmesh/core/test/result.py b/sqlmesh/core/test/result.py
index 5a8da783..11f168fb 100644
--- a/sqlmesh/core/test/result.py
+++ b/sqlmesh/core/test/result.py
@@ -1,10 +1,17 @@
 from __future__ import annotations
 
 import types
+import typing as t
 import unittest
 
 
 class ModelTextTestResult(unittest.TextTestResult):
+    successes: t.List[unittest.TestCase]
+
+    def __init__(self, *args: t.Any, **kwargs: t.Any):
+        super().__init__(*args, **kwargs)
+        self.successes = []
+
     def addFailure(
         self,
         test: unittest.TestCase,
@@ -22,3 +29,11 @@ class ModelTextTestResult(unittest.TextTestResult):
         """
         exctype, value, tb = err
         return super().addFailure(test, (exctype, value, None))  # type: ignore
+
+    def addSuccess(self, test: unittest.TestCase) -> None:
+        """Called when the test case test succeeds.
+
+        Args:
+            test: The test case
+        """
+        self.successes.append(test)
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 4e9d6e0f..ffa5eaaa 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1811,6 +1811,24 @@ def test_model_ctas_query():
         == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE ORDER BY 1 LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        WITH RECURSIVE a AS (
+            SELECT * FROM x
+        ), b AS (
+            SELECT * FROM a UNION ALL SELECT * FROM a
+        )
+        SELECT * FROM b
+
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'WITH RECURSIVE "a" AS (SELECT * FROM "x" AS "x" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 90c51588..035e71a1 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -597,7 +597,7 @@ def test_test_generation(tmp_path: Path) -> None:
     context.plan(auto_apply=True)
 
     input_queries = {
-        "sqlmesh_example.incremental_model": f"SELECT * FROM sqlmesh_example.incremental_model LIMIT 3"
+        "sqlmesh_example.incremental_model": "SELECT * FROM sqlmesh_example.incremental_model LIMIT 3"
     }
 
     with pytest.raises(ConfigError) as ex:
@@ -692,3 +692,11 @@ test_foo:
     ).run()
 
     _check_successful_or_raise(result)
+
+
+def test_successes(sushi_context: Context) -> None:
+    results = sushi_context.test()
+    successful_tests = [success.test_name for success in results.successes]  # type: ignore
+    assert len(successful_tests) == 2
+    assert "test_order_items" in successful_tests
+    assert "test_customer_revenue_by_day" in successful_tests