1010from pathlib import Path
1111from unittest .mock import patch
1212
13+ from rich .table import Table
14+ from rich .align import Align
15+
1316import numpy as np
1417import pandas as pd
1518from io import StringIO
2730from sqlmesh .utils .date import date_dict , pandas_timestamp_to_pydatetime , to_datetime
2831from sqlmesh .utils .errors import ConfigError , TestError
2932from sqlmesh .utils .yaml import load as yaml_load
33+ from sqlmesh .utils import Verbosity
3034
3135if t .TYPE_CHECKING :
3236 from sqlglot .dialects .dialect import DialectType
@@ -61,6 +65,8 @@ def __init__(
6165 preserve_fixtures : bool = False ,
6266 default_catalog : str | None = None ,
6367 concurrency : bool = False ,
68+ verbosity : Verbosity = Verbosity .DEFAULT ,
69+ rich_output : bool = True ,
6470 ) -> None :
6571 """ModelTest encapsulates a unit test for a model.
6672
@@ -84,6 +90,8 @@ def __init__(
8490 self .default_catalog = default_catalog
8591 self .dialect = dialect
8692 self .concurrency = concurrency
93+ self .verbosity = verbosity
94+ self .rich_output = rich_output
8795
8896 self ._fixture_table_cache : t .Dict [str , exp .Table ] = {}
8997 self ._normalized_column_name_cache : t .Dict [str , str ] = {}
@@ -278,6 +286,7 @@ def _to_hashable(x: t.Any) -> t.Any:
278286 check_like = True , # Ignore column order
279287 )
280288 except AssertionError as e :
289+ args : t .List [t .Any ] = []
281290 if expected .shape != actual .shape :
282291 _raise_if_unexpected_columns (expected .columns , actual .columns )
283292
@@ -291,10 +300,29 @@ def _to_hashable(x: t.Any) -> t.Any:
291300 if not unexpected_rows .empty :
292301 error_msg += f"\n \n Unexpected rows:\n \n { unexpected_rows } "
293302
294- e . args = (error_msg , )
303+ args . append (error_msg )
295304 else :
296- diff = expected .compare (actual ).rename (columns = {"self" : "exp" , "other" : "act" })
297- e .args = (f"Data mismatch (exp: expected, act: actual)\n \n { diff } " ,)
305+ diff = expected .compare (actual ).rename (
306+ columns = {"self" : "Expected" , "other" : "Actual" }
307+ )
308+
309+ if not self .rich_output :
310+ args .append (f"Data mismatch\n \n { diff } " )
311+ elif self .verbosity == Verbosity .DEFAULT :
312+ args .append (df_to_table ("Data mismatch" , diff ))
313+ else :
314+ from pandas import MultiIndex
315+
316+ levels = t .cast (MultiIndex , diff .columns ).levels [0 ]
317+ for col in levels :
318+ col_diff = diff [col ]
319+ if not col_diff .empty :
320+ table = df_to_table (
321+ f"[bold red]Column '{ col } ' mismatch[/bold red]" , col_diff
322+ )
323+ args .append (table )
324+
325+ e .args = (* args ,)
298326
299327 raise e
300328
@@ -316,6 +344,7 @@ def create_test(
316344 preserve_fixtures : bool = False ,
317345 default_catalog : str | None = None ,
318346 concurrency : bool = False ,
347+ verbosity : Verbosity = Verbosity .DEFAULT ,
319348 ) -> t .Optional [ModelTest ]:
320349 """Create a SqlModelTest or a PythonModelTest.
321350
@@ -361,6 +390,7 @@ def create_test(
361390 preserve_fixtures ,
362391 default_catalog ,
363392 concurrency ,
393+ verbosity ,
364394 )
365395 except Exception as e :
366396 raise TestError (f"Failed to create test { test_name } ({ path } )\n { str (e )} " )
@@ -676,6 +706,8 @@ def __init__(
676706 preserve_fixtures : bool = False ,
677707 default_catalog : str | None = None ,
678708 concurrency : bool = False ,
709+ verbosity : Verbosity = Verbosity .DEFAULT ,
710+ rich_output : bool = True ,
679711 ) -> None :
680712 """PythonModelTest encapsulates a unit test for a Python model.
681713
@@ -702,6 +734,8 @@ def __init__(
702734 preserve_fixtures ,
703735 default_catalog ,
704736 concurrency ,
737+ verbosity ,
738+ rich_output ,
705739 )
706740
707741 self .context = TestExecutionContext (
@@ -926,3 +960,41 @@ def _normalize_df_value(value: t.Any) -> t.Any:
926960 return {k : _normalize_df_value (v ) for k , v in zip (value ["key" ], value ["value" ])}
927961 return {k : _normalize_df_value (v ) for k , v in value .items ()}
928962 return value
963+
964+
965+ def df_to_table (
966+ header : str ,
967+ df : pd .DataFrame ,
968+ show_index : bool = True ,
969+ index_name : str = "Row" ,
970+ ) -> Table :
971+ """Convert a pandas.DataFrame obj into a rich.Table obj.
972+ Args:
973+ df (DataFrame): A Pandas DataFrame to be converted to a rich Table.
974+ rich_table (Table): A rich Table that should be populated by the DataFrame values.
975+ show_index (bool): Add a column with a row count to the table. Defaults to True.
976+ index_name (str, optional): The column name to give to the index column. Defaults to None, showing no value.
977+ Returns:
978+ Table: The rich Table instance passed, populated with the DataFrame values."""
979+
980+ rich_table = Table (title = f"[bold red]{ header } [/bold red]" , show_lines = True , min_width = 60 )
981+ if show_index :
982+ index_name = str (index_name ) if index_name else ""
983+ rich_table .add_column (index_name )
984+
985+ for column in df .columns :
986+ column_name = column if isinstance (column , str ) else ": " .join (str (col ) for col in column )
987+ if "expected" in column_name .lower ():
988+ column_name = f"[green]{ column_name } [/green]"
989+ else :
990+ column_name = f"[red]{ column_name } [/red]"
991+
992+ rich_table .add_column (Align .center (column_name ))
993+
994+ for index , value_list in enumerate (df .values .tolist ()):
995+ row = [str (index )] if show_index else []
996+ row += [str (x ) for x in value_list ]
997+ center = [Align .center (x ) for x in row ]
998+ rich_table .add_row (* center )
999+
1000+ return rich_table
0 commit comments