Refactor test_test to validate rich output [skip ci]

Author: VaggelisD

sqlmesh/core/console.py

@@ -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."""
 
 
@@ -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
 
 
@@ -1964,6 +1968,7 @@ def log_test_results(
         self, result: unittest.result.TestResult, output: t.Optional[str], target_dialect: str
     ) -> None:
         divider_length = 70
+
         if result.wasSuccessful():
             self._print("=" * divider_length)
             self._print(
@@ -2500,7 +2505,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 +2531,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 ''}"
@@ -3204,7 +3216,11 @@ 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")
+
+            self._print("```")
+            self._print(f"{output}\n\n")
+            self.log_unit_test_results(result)
+            self._print("```")
 
     def log_skipped_models(self, snapshot_names: t.Set[str]) -> None:
         if snapshot_names:

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

@@ -76,7 +76,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,10 @@ def _run_single_test(
 
     end_time = time.perf_counter()
 
+    from sqlmesh.core.console import get_console
+    console = get_console()
+    print(f"console {console}")
+
     combined_results.log_test_report(test_duration=end_time - start_time)
 
     return combined_results