Refactor test_test to validate rich output [skip ci]

Author: VaggelisD

sqlmesh/core/console.py

@@ -463,7 +463,7 @@ def plan(
 
     @abc.abstractmethod
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         """Display the test result and output.
 
@@ -498,7 +498,9 @@ def loading_stop(self, id: uuid.UUID) -> None:
         """Stop loading for the given id."""
 
     @abc.abstractmethod
-    def log_unit_test_results(self, result: ModelTextTestResult, test_duration: float) -> None:
+    def log_unit_test_results(
+        self, result: ModelTextTestResult, test_duration: t.Optional[float] = None
+    ) -> None:
         """Print the unit test results."""
 
 
@@ -674,7 +676,7 @@ def plan(
             plan_builder.apply()
 
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         pass
 
@@ -782,7 +784,9 @@ def start_destroy(self) -> bool:
     def stop_destroy(self, success: bool = True) -> None:
         pass
 
-    def log_unit_test_results(self, result: ModelTextTestResult, test_duration: float) -> None:
+    def log_unit_test_results(
+        self, result: ModelTextTestResult, test_duration: t.Optional[float] = None
+    ) -> None:
         pass
 
 
@@ -1961,9 +1965,13 @@ def _prompt_promote(self, plan_builder: PlanBuilder) -> None:
             plan_builder.apply()
 
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         divider_length = 70
+
+        self.log_unit_test_results(result)
+
+        self._print("\n")
         if result.wasSuccessful():
             self._print("=" * divider_length)
             self._print(
@@ -1980,7 +1988,7 @@ def log_test_results(
             )
             for test, _ in result.failures + result.errors:
                 if isinstance(test, ModelTest):
-                    self._print(f"Failure Test: {test.model.name} {test.test_name}")
+                    self._print(f"Failure Test: {test.path}::{test.test_name}")
             self._print("=" * divider_length)
             self._print(output)
 
@@ -2500,7 +2508,9 @@ def show_linter_violations(
         else:
             self.log_warning(msg)
 
-    def log_unit_test_results(self, result: ModelTextTestResult, test_duration: float) -> None:
+    def log_unit_test_results(
+        self, result: ModelTextTestResult, test_duration: t.Optional[float] = None
+    ) -> None:
         tests_run = result.testsRun
         errors = result.errors
         failures = result.original_failures
@@ -2524,22 +2534,27 @@ def log_unit_test_results(self, result: ModelTextTestResult, test_duration: floa
 
             if test_description := test_case.shortDescription():
                 self._print(test_description)
-            self._print(f"{unittest.TextTestResult.separator2}\n")
+            self._print(f"{unittest.TextTestResult.separator2}")
 
             if exception := failure[1]:
-                for arg in exception.args:
+                for i, arg in enumerate(exception.args):
+                    arg = f"Exception: {arg}" if isinstance(arg, str) else arg
                     self._print(arg)
-                    self._print("\n")
+
+                    if i < len(exception.args) - 1:
+                        self._print("\n")
 
         for test_case, error in errors:
             self._print(unittest.TextTestResult.separator1)
             self._print(f"ERROR: {test_case}")
+            self._print(f"{unittest.TextTestResult.separator2}")
             self._print(error)
 
         # Output final report
         self._print(unittest.TextTestResult.separator2)
+        test_duration_msg = f" in {test_duration:.3f}s" if test_duration else ""
         self._print(
-            f"Ran {tests_run} {'tests' if tests_run > 1 else 'test'} in {test_duration:.3f}s \n"
+            f"\nRan {tests_run} {'tests' if tests_run > 1 else 'test'}{test_duration_msg} \n"
         )
         self._print(
             f"{'OK' if is_success else 'FAILED'}{' (' + ', '.join(infos) + ')' if infos else ''}"
@@ -2817,7 +2832,7 @@ def radio_button_selected(change: t.Dict[str, t.Any]) -> None:
         self.display(radio)
 
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         import ipywidgets as widgets
 
@@ -3191,8 +3206,12 @@ def log_success(self, message: str) -> None:
         self._print(message)
 
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
+        # self._print("```")
+        self.log_unit_test_results(result)
+        # self._print("```\n\n")
+
         if result.wasSuccessful():
             self._print(
                 f"**Successfully Ran `{str(result.testsRun)}` Tests Against `{target_dialect}`**\n\n"
@@ -3204,6 +3223,7 @@ def log_test_results(
             for test, _ in result.failures + result.errors:
                 if isinstance(test, ModelTest):
                     self._print(f"* Failure Test: `{test.model.name}` - `{test.test_name}`\n\n")
+
             self._print(f"```{output}```\n\n")
 
     def log_skipped_models(self, snapshot_names: t.Set[str]) -> None:
@@ -3584,7 +3604,7 @@ def show_model_difference_summary(
             self._write(f"  Modified: {modified}")
 
     def log_test_results(
-        self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
+        self, result: ModelTextTestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         self._write("Test Results:", result)
 

sqlmesh/core/context.py

@@ -2053,7 +2053,7 @@ def test(
 
         test_meta = self.load_model_tests(tests=tests, patterns=match_patterns)
 
-        return run_tests(
+        result = run_tests(
             model_test_metadata=test_meta,
             models=self._models,
             config=self.config,
@@ -2066,6 +2066,10 @@ def test(
             default_catalog_dialect=self.config.dialect or "",
         )
 
+        self.console.log_test_results(result, output="", target_dialect=self.default_dialect)
+
+        return result
+
     @python_api_analytics
     def audit(
         self,

sqlmesh/core/test/definition.py

@@ -65,7 +65,6 @@ def __init__(
         default_catalog: str | None = None,
         concurrency: bool = False,
         verbosity: Verbosity = Verbosity.DEFAULT,
-        rich_output: bool = True,
     ) -> None:
         """ModelTest encapsulates a unit test for a model.
 
@@ -90,7 +89,6 @@ def __init__(
         self.dialect = dialect
         self.concurrency = concurrency
         self.verbosity = verbosity
-        self.rich_output = rich_output
 
         self._fixture_table_cache: t.Dict[str, exp.Table] = {}
         self._normalized_column_name_cache: t.Dict[str, str] = {}
@@ -142,6 +140,12 @@ def __init__(
 
         super().__init__()
 
+    def defaultTestResult(self) -> unittest.TestResult:
+        from sqlmesh.core.test.result import ModelTextTestResult
+        import sys
+
+        return ModelTextTestResult(stream=sys.stdout, descriptions=True, verbosity=self.verbosity)
+
     def shortDescription(self) -> t.Optional[str]:
         return self.body.get("description")
 
@@ -293,25 +297,23 @@ def _to_hashable(x: t.Any) -> t.Any:
             if expected.shape != actual.shape:
                 _raise_if_unexpected_columns(expected.columns, actual.columns)
 
-                error_msg = "Data mismatch (rows are different)"
+                args.append("Data mismatch (rows are different)")
 
                 missing_rows = _row_difference(expected, actual)
                 if not missing_rows.empty:
-                    error_msg += f"\n\nMissing rows:\n\n{missing_rows}"
+                    args.append(df_to_table("Missing rows", missing_rows))
 
                 unexpected_rows = _row_difference(actual, expected)
+
                 if not unexpected_rows.empty:
-                    error_msg += f"\n\nUnexpected rows:\n\n{unexpected_rows}"
+                    args.append(df_to_table("Unexpected rows", unexpected_rows))
 
-                args.append(error_msg)
             else:
                 diff = expected.compare(actual).rename(
                     columns={"self": "Expected", "other": "Actual"}
                 )
 
-                if not self.rich_output:
-                    args.append(f"Data mismatch\n\n{diff}")
-                elif self.verbosity == Verbosity.DEFAULT:
+                if self.verbosity == Verbosity.DEFAULT:
                     args.append(df_to_table("Data mismatch", diff))
                 else:
                     from pandas import MultiIndex
@@ -714,7 +716,6 @@ def __init__(
         default_catalog: str | None = None,
         concurrency: bool = False,
         verbosity: Verbosity = Verbosity.DEFAULT,
-        rich_output: bool = True,
     ) -> None:
         """PythonModelTest encapsulates a unit test for a Python model.
 
@@ -742,7 +743,6 @@ def __init__(
             default_catalog,
             concurrency,
             verbosity,
-            rich_output,
         )
 
         self.context = TestExecutionContext(
@@ -996,7 +996,7 @@ def df_to_table(
     rich_table = Table(title=f"[bold red]{header}[/bold red]", show_lines=True, min_width=60)
     if show_index:
         index_name = str(index_name) if index_name else ""
-        rich_table.add_column(index_name)
+        rich_table.add_column(Align.center(index_name))
 
     for column in df.columns:
         column_name = column if isinstance(column, str) else ": ".join(str(col) for col in column)

sqlmesh/core/test/result.py

@@ -19,6 +19,7 @@ def __init__(self, *args: t.Any, **kwargs: t.Any):
         self.successes = []
         self.original_failures: t.List[t.Tuple[unittest.TestCase, ErrorType]] = []
         self.original_errors: t.List[t.Tuple[unittest.TestCase, ErrorType]] = []
+        self.duration: t.Optional[float] = None
 
     def addSubTest(
         self,
@@ -76,7 +77,7 @@ def addSuccess(self, test: unittest.TestCase) -> None:
         super().addSuccess(test)
         self.successes.append(test)
 
-    def log_test_report(self, test_duration: float) -> None:
+    def log_test_report(self, test_duration: t.Optional[float] = None) -> None:
         """
         Log the test report following unittest's conventions.
 

sqlmesh/core/test/runner.py

@@ -184,6 +184,6 @@ def _run_single_test(
 
     end_time = time.perf_counter()
 
-    combined_results.log_test_report(test_duration=end_time - start_time)
+    combined_results.duration = end_time - start_time
 
     return combined_results