Use daff for diff formatting in unit testing

.changes/unreleased/Features-20231114-101555.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Use daff to render diff displayed in stdout when unit test fails
+time: 2023-11-14T10:15:55.689307-05:00
+custom:
+  Author: michelleark
+  Issue: "8558"

core/dbt/clients/agate_helper.py

@@ -3,6 +3,7 @@
 import agate
 import datetime
 import isodate
+import io
 import json
 import dbt.utils
 from typing import Iterable, List, Dict, Union, Optional, Any
@@ -137,6 +138,23 @@ def table_from_data_flat(data, column_names: Iterable[str]) -> agate.Table:
     )
 
 
+def json_rows_from_table(table: agate.Table) -> List[Dict[str, Any]]:
+    "Convert a table to a list of row dict objects"
+    output = io.StringIO()
+    table.to_json(path=output)  # type: ignore
+
+    return json.loads(output.getvalue())
+
+
+def list_rows_from_table(table: agate.Table) -> List[Any]:
+    "Convert a table to a list of lists, where the first element represents the header"
+    rows = [[col.name for col in table.columns]]
+    for row in table.rows:
+        rows.append(list(row.values()))
+
+    return rows
+
+
 def empty_table():
     "Returns an empty Agate table. To be used in place of None"
 

core/dbt/task/unit_test.py

@@ -1,13 +1,16 @@
+import agate
 from dataclasses import dataclass
 from dbt.dataclass_schema import dbtClassMixin
+import daff
 import threading
-from typing import Dict, Any, Optional, AbstractSet
-import io
+import re
+from typing import Dict, Any, Optional, AbstractSet, List
 
 from .compile import CompileRunner
 from .run import RunTask
 
 from dbt.adapters.factory import get_adapter
+from dbt.clients.agate_helper import list_rows_from_table, json_rows_from_table
 from dbt.contracts.graph.nodes import UnitTestNode
 from dbt.contracts.graph.manifest import Manifest
 from dbt.contracts.results import TestStatus, RunResult
@@ -25,16 +28,26 @@
 )
 from dbt.node_types import NodeType
 from dbt.parser.unit_tests import UnitTestManifestLoader
+from dbt.ui import green, red
+
+
+@dataclass
+class UnitTestDiff(dbtClassMixin):
+    actual: List[Dict[str, Any]]
+    expected: List[Dict[str, Any]]
+    rendered: str
 
 
 @dataclass
 class UnitTestResultData(dbtClassMixin):
     should_error: bool
     adapter_response: Dict[str, Any]
-    diff: Optional[str] = None
+    diff: Optional[UnitTestDiff] = None
 
 
 class UnitTestRunner(CompileRunner):
+    _ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
     def describe_node(self):
         return f"{self.node.resource_type} {self.node.name}"
 
@@ -95,15 +108,23 @@
         result = context["load_result"]("main")
         adapter_response = result["response"].to_dict(omit_none=True)
         table = result["table"]
-        actual = self._get_unit_test_table(table, "actual")
-        expected = self._get_unit_test_table(table, "expected")
-        should_error = actual.rows != expected.rows
-        diff = None
-        if should_error:
-            actual_output = self._agate_table_to_str(actual)
-            expected_output = self._agate_table_to_str(expected)
-
-            diff = f"\n\nActual:\n{actual_output}\n\nExpected:\n{expected_output}\n"
+        actual = self._get_unit_test_agate_table(table, "actual")
+        expected = self._get_unit_test_agate_table(table, "expected")
+
+        # generate diff, if exists
+        should_error, diff = False, None
+        daff_diff = self._get_daff_diff(expected, actual)
+        if daff_diff.hasDifference():
+            should_error = True
+            rendered = self._render_daff_diff(daff_diff)
+            rendered = f"\n\n{red('expected')} differs from {green('actual')}:\n\n{rendered}\n"
+
+            diff = UnitTestDiff(
+                actual=json_rows_from_table(actual),
+                expected=json_rows_from_table(expected),
+                rendered=rendered,
+            )
+
         return UnitTestResultData(
             diff=diff,
             should_error=should_error,
@@ -119,7 +140,7 @@
         failures = 0
         if result.should_error:
             status = TestStatus.Fail
-            message = result.diff
+            message = result.diff.rendered if result.diff else None
             failures = 1
 
         return RunResult(
@@ -136,22 +157,40 @@
     def after_execute(self, result):
         self.print_result_line(result)
 
-    def _get_unit_test_table(self, result_table, actual_or_expected: str):
+    def _get_unit_test_agate_table(self, result_table, actual_or_expected: str) -> agate.Table:
         unit_test_table = result_table.where(
             lambda row: row["actual_or_expected"] == actual_or_expected
         )
         columns = list(unit_test_table.columns.keys())
         columns.remove("actual_or_expected")
         return unit_test_table.select(columns)
 
-    def _agate_table_to_str(self, table) -> str:
-        # Hack to get Agate table output as string
-        output = io.StringIO()
-        if self.config.args.output == "json":
-            table.to_json(path=output)
-        else:
-            table.print_table(output=output, max_rows=None)
-        return output.getvalue().strip()
+    def _get_daff_diff(
+        self, expected: agate.Table, actual: agate.Table, ordered: bool = False
+    ) -> daff.TableDiff:
+
+        expected_daff_table = daff.PythonTableView(list_rows_from_table(expected))
+        actual_daff_table = daff.PythonTableView(list_rows_from_table(actual))
+
+        alignment = daff.Coopy.compareTables(expected_daff_table, actual_daff_table).align()
+        result = daff.PythonTableView([])
+
+        flags = daff.CompareFlags()
+        flags.ordered = ordered
+
+        diff = daff.TableDiff(alignment, flags)
+        diff.hilite(result)
+        return diff
+
+    def _render_daff_diff(self, daff_diff: daff.TableDiff) -> str:
+        result = daff.PythonTableView([])
+        daff_diff.hilite(result)
+        rendered = daff.TerminalDiffRender().render(result)
+        # strip colors if necessary
+        if not self.config.args.use_colors:
+            rendered = self._ANSI_ESCAPE.sub("", rendered)
+
+        return rendered
 
 
 class UnitTestSelector(ResourceTypeSelector):

core/setup.py

@@ -81,6 +81,7 @@
         "protobuf>=4.0.0",
         "pytz>=2015.7",
         "pyyaml>=6.0",
+        "daff>=1.3.46",
         "typing-extensions>=4.4",
         # ----
         # Match snowflake-connector-python, to ensure compatibility in dbt-snowflake