Add beam.testing.utils.equal_to function to display colorful differences

Author: Tomás Link

pyproject.toml

@@ -49,9 +49,9 @@ classifiers = [
   "Topic :: Scientific/Engineering",
 ]
 dependencies = [
+  "jinja2~=3.1",
   "pyyaml~=6.0",
   "rich~=14.0",
-  "jinja2~=3.1",
 ]
 
 [project.optional-dependencies]
@@ -182,7 +182,7 @@ strict = true
 ignore_missing_imports = true
 files = "src"
 mypy_path = "src"
-disable_error_code = ["union-attr", "no-any-return"]
+disable_error_code = ["union-attr", "no-any-return", "var-annotated"]
 explicit_package_bases = true
 
 [[tool.mypy.overrides]]

src/gfw/common/beam/testing/__init__.py

@@ -0,0 +1 @@
+"""Testing utilities for Apache Beam pipelines."""

src/gfw/common/beam/testing/utils.py

@@ -0,0 +1,95 @@
+"""General utilities for testing Apache Beam pipelines."""
+
+from itertools import zip_longest
+from typing import Any, Callable, Iterable, List, Sequence
+
+from apache_beam.testing.util import BeamAssertException
+from rich.console import Console, Group, RenderableType
+
+from gfw.common.diff import compare_items, render_diff_panel
+from gfw.common.sorting import sort_dicts
+
+
+def _default_equals_fn(e: Any, a: Any) -> bool:
+    return e == a
+
+
+def _raise_with_diff(diffs: Sequence[RenderableType]) -> None:
+    # Set up a Rich Console that records output
+    console = Console(record=True, force_terminal=True, width=130)
+
+    # Render diffs to console (only into memory, not to screen)
+    console.print(Group(*diffs))
+
+    # Export the captured diff as text with ANSI codes
+    diff_text = console.export_text(styles=True)
+
+    # Raise exception with embedded colored diff
+    raise BeamAssertException(f"PCollection contents differ: \n{diff_text}.")
+
+
+def equal_to(
+    expected: List[Any], equals_fn: Callable[[Any, Any], bool] = _default_equals_fn
+) -> Callable[[List[Any]], None]:
+    """Drop-in replacement for `apache_beam.testing.util.equal_to` with rich diff output.
+
+    This matcher performs unordered comparison of top-level elements in actual and expected
+    PCollection outputs, just like Apache Beam's `equal_to`. However, it adds a rich diff
+    visualization to help debug mismatches by rendering side-by-side differences.
+
+    Use in tests with `assert_that(pcoll, equal_to(expected))`.
+
+    Note:
+        - Only top-level permutations are considered equal:
+          `[1, 2]` and `[2, 1]` are equal, but `[[1, 2]]` and `[[2, 1]]` are not.
+
+        - If elements are not directly comparable, a fallback comparison using
+          a custom equality function or deep diff is used. This helps handle:
+          1) Collections with types that don't have a deterministic sort order
+             (e.g., pyarrow Tables as of 0.14.1).
+          2) Collections containing elements of different types.
+
+    Args:
+        expected: Iterable of expected PCollection elements.
+        equals_fn: Optional function `(expected_item, actual_item) -> bool` to customize equality.
+
+    Returns:
+        A matcher function for use with `apache_beam.testing.util.assert_that`.
+    """
+
+    def _matcher(actual: Iterable[Any]) -> None:
+        expected_list = [sort_dicts(e) for e in expected]
+        actual_list = [sort_dicts(e) for e in actual]
+
+        try:
+            if actual_list == expected_list:
+                return
+        except TypeError:
+            pass
+
+        # Slower method, fallback comparison.
+        unmatched_expected = expected_list[:]
+        unmatched_actual = []
+        for a in actual_list:
+            for i, e in enumerate(unmatched_expected):
+                if equals_fn(e, a):
+                    unmatched_expected.pop(i)
+                    break
+            else:
+                unmatched_actual.append(a)
+
+        if not unmatched_actual and not unmatched_expected:
+            return
+
+        diffs = []
+        for i, (a, b) in enumerate(
+            zip_longest(unmatched_actual, unmatched_expected, fillvalue={}), 1
+        ):
+            left, right, changed = compare_items(a, b)
+            if changed:
+                diffs.append(render_diff_panel(left, right, i))
+
+        if diffs:  # Diffs found. Raise exception with colorized output.
+            _raise_with_diff(diffs)
+
+    return _matcher

src/gfw/common/diff.py

@@ -0,0 +1,87 @@
+"""General utilities for generating diffs between objects."""
+
+from difflib import ndiff
+from typing import Any, Tuple
+
+from rich.columns import Columns
+from rich.panel import Panel
+from rich.pretty import pretty_repr
+
+
+def diff_lines(a: str, b: str) -> Tuple[str, str, bool]:
+    """Generate a line-by-line diff of two strings with rich markup.
+
+    Args:
+        a:
+            First multi-line string to compare.
+        b:
+            Second multi-line string to compare.
+
+    Returns:
+        A tuple of (a_diff, b_diff, changed) where:
+        - a_diff: The first string annotated with diff highlights.
+        - b_diff: The second string annotated with diff highlights.
+        - changed: True if any differences were found, False otherwise.
+    """
+    a_lines, b_lines = a.splitlines(), b.splitlines()
+    a_out, b_out = [], []
+    changed = False
+    for line in ndiff(a_lines, b_lines):
+        tag, content = line[0], line[2:]
+        if tag == " ":
+            a_out.append(f"  {content}")
+            b_out.append(f"  {content}")
+        elif tag == "-":
+            changed = True
+            a_out.append(f"[red]- {content}[/red]")
+            b_out.append("")  # line not in b
+        elif tag == "+":
+            changed = True
+            a_out.append("")  # line not in a
+            b_out.append(f"[green]+ {content}[/green]")
+
+    return "\n".join(a_out), "\n".join(b_out), changed
+
+
+def compare_items(a: Any, b: Any) -> Tuple[str, str, bool]:
+    """Generate a rich diff of two objects' pretty-printed representations.
+
+    Args:
+        a:
+            First object to compare.
+        b:
+            Second object to compare.
+
+    Returns:
+        The object returned by diff_lines.
+    """
+    return diff_lines(
+        pretty_repr(a, indent_size=4, max_width=20),
+        pretty_repr(b, indent_size=4, max_width=20),
+    )
+
+
+def render_diff_panel(left: str, right: str, idx: int) -> Columns:
+    """Render side-by-side panels of diff strings for visual comparison.
+
+    Args:
+        left:
+            The left-side diff string (usually actual output).
+
+        right:
+            The right-side diff string (usually expected output).
+
+        idx:
+            Index number for labeling the diff panels.
+
+    Returns:
+        A rich Columns object containing two Panels side-by-side.
+    """
+    return Columns(
+        [
+            Panel(left, title=f"Actual #{idx}", expand=True),
+            Panel(right, title=f"Expected #{idx}", expand=True),
+        ],
+        expand=True,
+        equal=True,
+    )

src/gfw/common/sorting.py

@@ -0,0 +1,29 @@
+"""Utilities for sorting data structures."""
+
+from typing import Any
+
+
+def sort_dicts(obj: Any) -> Any:
+    """Recursively sorts dict keys to get consistent ordering for comparison.
+
+    Lists, tuples, and other types are returned unchanged (except their contents
+    get sorted recursively if they are dicts).
+
+    Args:
+        obj: Any nested structure (dict, list, tuple, or other).
+
+    Returns:
+        A new structure with dict keys sorted recursively.
+    """
+    if isinstance(obj, dict):
+        # Sort keys and recursively apply to values
+        return {k: sort_dicts(obj[k]) for k in sorted(obj)}
+    elif isinstance(obj, list):
+        # Recursively apply to each element
+        return [sort_dicts(e) for e in obj]
+    elif isinstance(obj, tuple):
+        # Recursively apply to each element and keep tuple type
+        return tuple(sort_dicts(e) for e in obj)
+    else:
+        # Return base case unchanged
+        return obj

tests/beam/testing/__init__.py

@@ -0,0 +1,104 @@
+import pytest
+from apache_beam.testing.util import BeamAssertException
+
+from gfw.common.beam.testing.utils import equal_to, _default_equals_fn
+
+
+def test_default_equals_fn():
+    assert _default_equals_fn(1, 1)
+    assert not _default_equals_fn(1, 2)
+
+
+def test_equal_to_match_exact_order():
+    expected = [1, 2, 3]
+    actual = [1, 2, 3]
+
+    matcher = equal_to(expected)
+    matcher(actual)  # should not raise
+
+
+def test_equal_to_match_different_order():
+    expected = [1, 2, 3]
+    actual = [3, 1, 2]
+
+    matcher = equal_to(expected)
+    matcher(actual)  # should not raise
+
+
+def test_equal_to_empty_lists():
+    matcher = equal_to([])
+    matcher([])  # should not raise
+
+
+def test_equal_to_mismatch_raises():
+    expected = [1, 2]
+    actual = [1, 3]
+
+    matcher = equal_to(expected)
+
+    with pytest.raises(BeamAssertException) as e:
+        matcher(actual)
+
+    # The exception message should contain a substring hinting at mismatch
+    assert "PCollection contents differ" in str(e.value)
+
+
+def test_equal_to_type_error_handling():
+    class Uncomparable:
+        def __eq__(self, other):
+            raise TypeError()
+
+    a = [Uncomparable()]
+    b = [Uncomparable()]
+
+    def safe_equals(x, y):
+        return type(x) is type(y)
+
+    matcher = equal_to(b, equals_fn=safe_equals)
+
+    # Should not raise, because our fallback considers them equal by type
+    matcher(a)
+
+
+def test_equal_to_custom_equals_fn():
+    expected = [1, 2, 3]
+    actual = [3, 2, 1]
+
+    def reversed_equals(e, a):
+        return e == a
+
+    matcher = equal_to(expected, equals_fn=reversed_equals)
+    matcher(actual)  # should not raise
+
+    # A custom equals that never matches causes exception
+    def never_equals(e, a):
+        return False
+
+    matcher = equal_to(expected, equals_fn=never_equals)
+    with pytest.raises(BeamAssertException):
+        matcher(actual)
+
+
+def test_equal_to_handles_nested_dicts_order():
+    expected = [{"b": 1, "a": 2}]
+    actual = [{"a": 2, "b": 1}]
+
+    matcher = equal_to(expected)
+    matcher(actual)  # Should not raise because dict keys sorted recursively
+
+
+def test_equal_to_handles_unmatched_extra_and_missing():
+    expected = [1, 2]
+    actual = [1, 2, 3]
+
+    matcher = equal_to(expected)
+
+    with pytest.raises(BeamAssertException):
+        matcher(actual)
+
+    expected = [1, 2, 3]
+    actual = [1, 2]
+
+    matcher = equal_to(expected)
+    with pytest.raises(BeamAssertException):
+        matcher(actual)

tests/beam/testing/test_utils.py

@@ -37,7 +37,10 @@ def test_read_and_decode_from_pubsub():
     pubsub_messages = [
         dict(
             data=b'{"test": 123}',
-            attributes={"key": "value"},
+            attributes={
+                "key2": "value2",
+                "key1": "value1",
+            },
         )
     ]
 
@@ -57,7 +60,10 @@ def test_read_and_decode_from_pubsub():
         expected = [
             {
                 "data": '{"test": 123}',
-                "attributes": {"key": "value"}
+                "attributes": {
+                    "key1": "value1",
+                    "key2": "value2",
+                }
             }
         ]