From a8b9fa84ee44632bfb4a4329c3e95f93f7a7f90a Mon Sep 17 00:00:00 2001 From: Olivier 'reivilibre Date: Mon, 28 Sep 2026 10:23:27 +0200 Subject: [PATCH] Add utilities for sharded stream tokens Extracted from #20186 (branch rei/se_sharded_tok at 8e10c4def4) so that receipt paging for application services can use them without waiting for the sticky events changes in that PR. This is that branch's net change to these files relative to develop: the helpers in id_generators.py for bounding a query on a multi-writer stream by sharded tokens and for advancing a sharded token after a limited read, their tests, and the test assertion helpers those tests use. --- synapse/storage/util/id_generators.py | 199 ++++++++++++++- tests/storage/test_id_generators.py | 339 +++++++++++++++++++++++++- tests/unittest.py | 167 ++++++++++++- 3 files changed, 702 insertions(+), 3 deletions(-) diff --git a/synapse/storage/util/id_generators.py b/synapse/storage/util/id_generators.py index fbf75d5a2a..d370f2f815 100644 --- a/synapse/storage/util/id_generators.py +++ b/synapse/storage/util/id_generators.py @@ -36,6 +36,7 @@ from typing import ( ) import attr +from immutabledict import immutabledict from prometheus_client import Gauge from sortedcontainers import SortedList, SortedSet @@ -48,9 +49,10 @@ from synapse.storage.database import ( LoggingTransaction, make_in_list_sql_clause, ) -from synapse.storage.engines import PostgresEngine +from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine from synapse.storage.types import Cursor from synapse.storage.util.sequence import build_sequence_generator +from synapse.types import MultiWriterStreamToken if TYPE_CHECKING: from synapse.notifier import ReplicationNotifier @@ -1019,3 +1021,198 @@ class _MultiWriterCtxManager: ) return False + + +def make_multiwriter_sharded_token_bounds_sql( + db_engine: BaseDatabaseEngine, + *, + stream_id_column: str, + instance_name_column: str, + from_token_exclusive: MultiWriterStreamToken, + to_token_inclusive: MultiWriterStreamToken, +) -> tuple[str, Sequence[str | int]]: + """ + Build an SQL clause that, in a multi-writer stream table, + matches the rows within the bounds of the `MultiWriterStreamToken` tokens provided. + + Bounds: from_token_exclusive < ... <= to_token_inclusive + + Arguments: + stream_id_column: the name of the `stream_id` column + Should be prefixed with the table name or alias + when used in a multi-table SELECT statement. + instance_name_column: the name of the `instance_name` column. + Should be prefixed with the table name or alias + when used in a multi-table SELECT statement. + from_token_exclusive: + MultiWriterStreamToken representing the highest position that has 'already been seen'. + Exclusive lower bound + to_token_inclusive: + MultiWriterStreamToken representing the highest position that has visible data. + Inclusive upper bound + + Example: + + make_multiwriter_sharded_token_bounds_sql( + stream_id_column="se.stream_id", + instance_name_column="se.instance_name", + # m5~1.8 + from_token_exclusive=MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 8}) + ), + # m10~2.14 + to_token_inclusive=MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ) + + gives: + + ( + 5 < se.stream_id AND se.stream_id <= 14 + AND NOT (se.instance_name = 'worker1' AND se.stream_id <= 8) + AND ( + se.stream_id <= 10 + OR (se.instance_name = 'worker2' AND se.stream_id <= 14) + ) + ) + + which matches, for example: + - a `worker1` row at position 9 (past the 8 we had already read from + `worker1`, and within the 10 that is visible for every writer); + - a `worker2` row at position 12 (past 5, and `worker2` is + explicitly visible up to 14); + but not: + - a `worker1` row at position 7, which the `from` token says we have + already read; + - a `worker3` row at position 12, which is beyond the position that + is visible for writers other than `worker2`. + + See also: + - `advance_multiwriter_sharded_token_after_partial_read` to create the next + `from_token_exclusive` after the result of reading a limited set of rows. + """ + + is_not_distinct_from = "IS NOT DISTINCT FROM" + if isinstance(db_engine, Sqlite3Engine): + # TODO(SQLite 3.39.0): Drop this compatibility code + # SQL standard `IS NOT DISTINCT FROM` was not introduced until 3.39.0 + # https://sqlite.org/releaselog/3_39_0.html + is_not_distinct_from = "IS" + + # The SQL we build will fundamentally consist of many clauses ANDed together. + # + # We start with an envelope delimited by the two outermost, writer-independent, bounds. + # We know everything before the lowest position of `from` and everything after the highest + # position of `to` must be out-of-bounds. + # This pair of bounds is the most likely to be index-friendly. + clauses = [f"? < {stream_id_column}", f"{stream_id_column} <= ?"] + values: list[int | str] = [ + from_token_exclusive.stream, + to_token_inclusive.get_max_stream_pos(), + ] + + # Now for every writer where we had already read further ahead than the baseline (lowest) position of `from`, + # we whittle away what we have already seen from the lower end of the envelope. + for instance_name, pos in from_token_exclusive.instance_map.items(): + # We need `IS NOT DISTINCT FROM` (analogous to `=` but treats `NULL` as a known value) + # for legacy rows where `instance_name` is `NULL`, such as before the stream was sharded. + clauses.append( + f"NOT ({instance_name_column} {is_not_distinct_from} ? AND {stream_id_column} <= ?)" + ) + values.extend((instance_name, pos)) + + # Now we need to restrict the upper end of the envelope to only include those rows where either: + # - the row is in a region where all writers' rows are visible; or + upper_or_clauses = [f"{stream_id_column} <= ?"] + upper_or_values: list[int | str] = [to_token_inclusive.stream] + # - the row's writer is explicitly visible ahead of the baseline (minimum) position + for instance_name, pos in to_token_inclusive.instance_map.items(): + upper_or_clauses.append( + f"({instance_name_column} {is_not_distinct_from} ? AND {stream_id_column} <= ?)" + ) + upper_or_values.extend((instance_name, pos)) + + clauses.append("( " + " OR ".join(upper_or_clauses) + " )") + values.extend(upper_or_values) + + return "( " + " AND ".join(clauses) + " )", values + + +def advance_multiwriter_sharded_token_after_partial_read( + *, + from_token_exclusive: MultiWriterStreamToken, + to_token_inclusive: MultiWriterStreamToken, + last_read_stream_id: int, +) -> MultiWriterStreamToken: + """ + Calculates the 'read up to' token after reading a limited number of rows + between two `MultiWriterStreamToken` bounds. + Pairs with `make_multiwriter_sharded_token_bounds_sql`. + + The read operation MUST have been ordered by the stream ID, e.g. + using `ORDER BY stream_id`. + + Caution: + For streams where one fact is represented by multiple rows with + the same `stream_id`, the ENTIRE fact corresponding to `last_read_stream_id` + MUST have been processed because this function won't help you 'pause' + in the middle of a fact. + + Arguments: + from_token_exclusive: the exclusive lower bound the rows were read with. + to_token_inclusive: the inclusive upper bound the rows were read with. + last_read_stream_id: the `stream_id` of the last row that was read. + + Returns: + A token that could be used as the `from_token_exclusive` for a subsequent + read, in order to continue reading rows in range. + + The token represents the maximum 'read up to' position. + Rows with positions strictly above the token have not yet been read. + """ + assert from_token_exclusive.stream < last_read_stream_id, "read was out of bounds" + + # Calculate the new baseline position that we have read up to, + # which applies across all workers. + new_baseline = max( + # Must be at least as far as it was before (doesn't go backwards) + from_token_exclusive.stream, + min( + # Simple case: this is just `last_read_stream_id`, + # the `stream_id` of the last row we read. + last_read_stream_id, + # Complicated case: some writers may not have advanced to `last_read_stream_id` yet, + # so we have to cap off at the baseline 'to' position (`to_token_inclusive.stream`) + to_token_inclusive.stream, + ), + ) + + # Now consider whether any workers are ahead of the baseline. + # We do this by calculating the position of each relevant worker. + # (Relevant by virtue of being in either the `from` or `to` token.) + advanced_instance_map = {} + for instance_name in ( + from_token_exclusive.instance_map.keys() + | to_token_inclusive.instance_map.keys() + ): + # Calculate the worker's position + worker_pos = max( + # Must be at least as far as it was before (doesn't go backwards) + from_token_exclusive.get_stream_pos_for_instance(instance_name), + # Then: + # - up to `last_read_stream_id` + # - unless its rows aren't visible yet, so cap off at the `to` position + # Either one of these cases will be at least as high as `new_baseline`. + min( + last_read_stream_id, + to_token_inclusive.get_stream_pos_for_instance(instance_name), + ), + ) + + if worker_pos > new_baseline: + advanced_instance_map[instance_name] = worker_pos + + return MultiWriterStreamToken( + stream=new_baseline, instance_map=immutabledict(advanced_instance_map) + ) diff --git a/tests/storage/test_id_generators.py b/tests/storage/test_id_generators.py index 0666eec00c..8eb2befe9c 100644 --- a/tests/storage/test_id_generators.py +++ b/tests/storage/test_id_generators.py @@ -19,8 +19,11 @@ # # +import sqlite3 from unittest import mock +from immutabledict import immutabledict + from twisted.internet.defer import CancelledError, Deferred, ensureDeferred from twisted.internet.testing import MemoryReactor @@ -31,9 +34,12 @@ from synapse.storage.database import ( LoggingDatabaseConnection, LoggingTransaction, ) +from synapse.storage.engines import Sqlite3Engine from synapse.storage.types import Cursor from synapse.storage.util.id_generators import ( MultiWriterIdGenerator, + advance_multiwriter_sharded_token_after_partial_read, + make_multiwriter_sharded_token_bounds_sql, stream_current_position_gauge, ) from synapse.storage.util.sequence import ( @@ -41,9 +47,10 @@ from synapse.storage.util.sequence import ( PostgresSequenceGenerator, SequenceGenerator, ) +from synapse.types import MultiWriterStreamToken from synapse.util.clock import Clock -from tests.unittest import HomeserverTestCase +from tests.unittest import HomeserverTestCase, TestCase from tests.utils import USE_POSTGRES_FOR_TESTS @@ -915,3 +922,333 @@ class MultiTableMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7) self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7) self.assertEqual(second_id_gen.get_persisted_upto_position(), 7) + + +class ShardedTokenHelpersPureTestCase(TestCase): + """ + Non-database tests for the helpers for reading multi-writer streams. + """ + + def test_bounds_sql_documented_example(self) -> None: + """ + Tests that the example in the docstring is what we actually generate. + """ + clause, values = make_multiwriter_sharded_token_bounds_sql( + Sqlite3Engine({}), + stream_id_column="se.stream_id", + instance_name_column="se.instance_name", + from_token_exclusive=MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 8}) + ), + to_token_inclusive=MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ) + self.assertEqualNormalisingWhitespace( + clause, + """ + ( + ? < se.stream_id + AND se.stream_id <= ? + AND NOT (se.instance_name IS ? AND se.stream_id <= ?) + AND ( + se.stream_id <= ? + OR (se.instance_name IS ? AND se.stream_id <= ?) + ) + ) + """, + ) + self.assertEqual(list(values), [5, 14, "worker1", 8, 10, "worker2", 14]) + + def test_token_after_partial_read_does_not_go_backwards(self) -> None: + """ + Tests that advancing the token after a partial read doesn't let it go + backwards. + + This is relevant because Synapse workers don't always advance their current + position at the same time. + """ + # As a scenario: the client has already read up to a baseline position of 60, + # but this reader worker has only caught up to 10. + # It can nonetheless see that worker2 has reached 70. + from_token = MultiWriterStreamToken(stream=60) + to_token = MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 70}) + ) + + resume_token = advance_multiwriter_sharded_token_after_partial_read( + from_token_exclusive=from_token, + to_token_inclusive=to_token, + last_read_stream_id=64, + ) + + # worker2 advances to what we read; everyone else stays where they were. + self.assertEqual( + resume_token, + MultiWriterStreamToken( + stream=60, instance_map=immutabledict({"worker2": 64}) + ), + ) + self.assertTrue( + from_token.is_before_or_eq(resume_token), + f"Expected {from_token} <= {resume_token}", + ) + + +class ShardedTokenHelpersDatabaseTestCase(TestCase): + """Tests for the helpers that read a range of a multi-writer stream. + + These don't need a homeserver: they only exercise the SQL that + `make_multiwriter_sharded_token_bounds_sql` builds, against a throwaway + SQLite table. + + NOTE: `ShardedTokenHelpersDatabaseNullTestCase` inherits all these tests + with a separate set of tables. + """ + + ROWS: list[tuple[int, str | None]] = [ + # (stream_id, instance_name) + (6, "worker1"), + (7, "worker1"), + (8, "worker2"), + (9, "worker1"), + (10, "worker3"), + (11, "worker1"), + (12, "worker2"), + (13, "worker3"), + (14, "worker2"), + ] + """ + (stream_id, instance_name) rows to insert into the example stream table. + """ + + SHARDED_TOKEN_RANGES = [ + ( + MultiWriterStreamToken(stream=5), + MultiWriterStreamToken(stream=14), + ), + ( + MultiWriterStreamToken(stream=5), + MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 9}) + ), + MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 7}) + ), + MultiWriterStreamToken( + stream=8, instance_map=immutabledict({"worker1": 11, "worker3": 13}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 9, "worker2": 8}) + ), + MultiWriterStreamToken( + stream=10, + instance_map=immutabledict( + {"worker1": 11, "worker2": 14, "worker3": 13} + ), + ), + ), + ] + """(from, to) token pairs to exercise the bounds against `ROWS`.""" + + def setUp(self) -> None: + # Just used for detecting SQL language support, + # doesn't actually connect... + self.database_engine = Sqlite3Engine({}) + self.conn = sqlite3.connect(":memory:") + self.addCleanup(self.conn.close) + self.conn.execute( + "CREATE TABLE streamtable (stream_id INTEGER PRIMARY KEY, instance_name TEXT)" + ) + self.conn.executemany("INSERT INTO streamtable VALUES (?, ?)", self.ROWS) + + def _select( + self, + from_token: MultiWriterStreamToken, + to_token: MultiWriterStreamToken, + limit: int | None = None, + ) -> list[tuple[int, str]]: + """ + Helper that selects the rows within the given bounds, in stream order, + using `make_multiwriter_sharded_token_bounds_sql`. + + Bounds: from_token < ... <= to_token + """ + clause, values = make_multiwriter_sharded_token_bounds_sql( + self.database_engine, + stream_id_column="streamtable.stream_id", + instance_name_column="streamtable.instance_name", + from_token_exclusive=from_token, + to_token_inclusive=to_token, + ) + + limit_clause = "" + if limit is not None: + limit_clause = f"LIMIT {limit}" + + return list( + self.conn.execute( + f""" + SELECT stream_id, instance_name + FROM streamtable + WHERE {clause} + ORDER BY stream_id ASC + {limit_clause} + """, + values, + ) + ) + + def _expected( + self, + from_token: MultiWriterStreamToken, + to_token: MultiWriterStreamToken, + ) -> list[tuple[int, str | None]]: + """ + Helper that returns the rows that ought to be within the given bounds. + It uses `MultiWriterStreamToken.is_stream_position_in_range` as the 'ground truth'. + + Bounds: from_token < ... <= to_token + """ + + return [ + (stream_id, instance_name) + for stream_id, instance_name in self.ROWS + if MultiWriterStreamToken.is_stream_position_in_range( + from_token, to_token, instance_name, stream_id + ) + ] + + def test_bounds_sql_selects_expected_rows(self) -> None: + """ + Tests that the `make_multiwriter_sharded_token_bounds_sql` selects the correct rows, + using `MultiWriterStreamToken.is_stream_position_in_range` as the 'ground truth'. + """ + for from_token, to_token in self.SHARDED_TOKEN_RANGES: + with self.subTest(from_token=str(from_token), to_token=str(to_token)): + self.assertEqual( + self._select(from_token, to_token), + self._expected(from_token, to_token), + ) + + def test_token_after_partial_read(self) -> None: + """ + Tests that when advancing the token after a partial read, the subsequent read returns all the remaining rows + but no more (especially not duplicates). + """ + for from_token, to_token in self.SHARDED_TOKEN_RANGES: + all_rows = self._expected(from_token, to_token) + for limit in range(1, len(all_rows) + 1): + with self.subTest( + from_token=str(from_token), to_token=str(to_token), limit=limit + ): + read = self._select(from_token, to_token, limit=limit) + resume_token = advance_multiwriter_sharded_token_after_partial_read( + from_token_exclusive=from_token, + to_token_inclusive=to_token, + last_read_stream_id=read[-1][0], + ) + + # The resumption point must be somewhere within the range we + # were asked to read. + self.assertTrue( + from_token.is_before_or_eq(resume_token), + f"Expected {from_token} <= {resume_token}", + ) + self.assertTrue( + resume_token.is_before_or_eq(to_token), + f"Expected {resume_token} <= {to_token}", + ) + + self.assertEqual( + read, + self._expected(from_token, resume_token), + f"The rows we read between {from_token} < ... <= {resume_token} seem wrong.", + ) + + self.assertEqual( + read + self._select(resume_token, to_token), + all_rows, + f"We did a limited read of {limit} rows within {from_token} < ... <= {to_token}, " + f"giving us an advanced token {resume_token}. " + f"We then read {resume_token} < ... <= {to_token} and expected to get " + "all the rows within range (in order, no duplicates), but didn't.", + ) + + +class ShardedTokenHelpersDatabaseNullTestCase(ShardedTokenHelpersDatabaseTestCase): + """ + Variant of `ShardedTokenHelpersDatabaseTestCase` where there are rows with + `instance_name` being `NULL`, corresponding to rows that existed before the stream was sharded. + """ + + ROWS = [ + # (stream_id, instance_name) + (4, None), + (5, None), + (6, None), + (7, None), + (8, "worker2"), + (9, "worker1"), + (10, "worker3"), + (11, "worker1"), + (12, "worker2"), + (13, "worker3"), + (14, "worker2"), + ] + """ + (stream_id, instance_name) rows to insert into the example stream table. + """ + + SHARDED_TOKEN_RANGES = [ + ( + MultiWriterStreamToken(stream=5), + MultiWriterStreamToken(stream=14), + ), + ( + MultiWriterStreamToken(stream=5), + MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 9}) + ), + MultiWriterStreamToken( + stream=10, instance_map=immutabledict({"worker2": 14}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 7}) + ), + MultiWriterStreamToken( + stream=8, instance_map=immutabledict({"worker1": 11, "worker3": 13}) + ), + ), + ( + MultiWriterStreamToken( + stream=5, instance_map=immutabledict({"worker1": 9, "worker2": 8}) + ), + MultiWriterStreamToken( + stream=10, + instance_map=immutabledict( + {"worker1": 11, "worker2": 14, "worker3": 13} + ), + ), + ), + ] + """(from, to) token pairs to exercise the bounds against `ROWS`.""" diff --git a/tests/unittest.py b/tests/unittest.py index f3dc91c58a..b67a0d9f2c 100644 --- a/tests/unittest.py +++ b/tests/unittest.py @@ -20,6 +20,7 @@ # # import copy +import difflib import functools import gc import hashlib @@ -27,9 +28,11 @@ import hmac import json import logging import os +import re import secrets import sys import time +from collections import Counter from collections.abc import Set from typing import ( TYPE_CHECKING, @@ -43,6 +46,7 @@ from typing import ( NoReturn, Optional, Protocol, + Sequence, TypeVar, ) from unittest.mock import Mock, patch @@ -127,6 +131,15 @@ def _use_colour() -> bool: return sys.stdout.isatty() +def _is_hashable(item: object) -> bool: + """Whether `item` can be hashed.""" + try: + hash(item) + except TypeError: + return False + return True + + class _TypedFailure(Generic[_ExcType], Protocol): """Extension to twisted.Failure, where the 'value' has a certain type.""" @@ -398,7 +411,119 @@ class TestCase(unittest.TestCase): extra_message = "" if message is not None: - extra_message = "f\n{message}" + extra_message = f"\n{message}" + + self.fail(f"{diff_message}{extra_message}") + + def _fail_with_sequence_inequality( + self, + actual_items: Sequence[TV], + expected_items: Sequence[TV], + message: str | None, + ) -> Never: + """ + Fail the current test, printing a rich message showing the sequence + inequality. + + As sequences are ordered, we align them with a sequence matcher, so that a single inserted + or deleted item doesn't make every subsequent item look wrong. + An item which appears in both lists but at a different position is marked as 'moved'. + + `actual_items` and `expected_items` must contain only hashable items. + """ + if _use_colour(): + BOLD = "\033[1m" + DIM_GREY = "\033[2;37m" + RED = "\033[31m" + BRIGHT_YELLOW = "\033[93m" + RESET = "\033[0m" + else: + BOLD = DIM_GREY = RED = BRIGHT_YELLOW = RESET = "" + + # If it's correct, use dim grey: we don't want to draw your attention to it + CORRECT_MARKER = f"{DIM_GREY} ok" + + # If it's wrong, use red: that's the most important information here + MISSING_MARKER = f" {RED}missing{RESET}{BRIGHT_YELLOW}" + UNEXPECTED_MARKER = f"{RED}unwanted{RESET}{BRIGHT_YELLOW}" + + # A 'moved' item is present in both lists, but at a different position + MOVED_MARKER = f"{BRIGHT_YELLOW} moved{RESET}{BRIGHT_YELLOW}" + + # Match up items that appear in the same order in both lists. This means + # a lone insertion/deletion doesn't shift everything after it. + matcher = difflib.SequenceMatcher( + a=expected_items, b=actual_items, autojunk=False + ) + expected_markers = [MISSING_MARKER] * len(expected_items) + actual_markers = [UNEXPECTED_MARKER] * len(actual_items) + for block in matcher.get_matching_blocks(): + for offset in range(block.size): + expected_markers[block.a + offset] = CORRECT_MARKER + actual_markers[block.b + offset] = CORRECT_MARKER + + # Anything left over which is present on both sides (by value) is + # 'moved' rather than missing/unexpected. We only move as many items as + # appear in both lists, so duplicates aren't over-counted. + unmatched_expected = Counter( + expected_items[i] + for i, marker in enumerate(expected_markers) + if marker != CORRECT_MARKER + ) + unmatched_actual = Counter( + actual_items[i] + for i, marker in enumerate(actual_markers) + if marker != CORRECT_MARKER + ) + movable = unmatched_expected & unmatched_actual + expected_budget = movable.copy() + actual_budget = movable.copy() + for i, item in enumerate(expected_items): + if expected_markers[i] != CORRECT_MARKER and expected_budget[item] > 0: + expected_markers[i] = MOVED_MARKER + expected_budget[item] -= 1 + for i, item in enumerate(actual_items): + if actual_markers[i] != CORRECT_MARKER and actual_budget[item] > 0: + actual_markers[i] = MOVED_MARKER + actual_budget[item] -= 1 + + used_markers = set() + expected_lines: list[str] = [] + for index, (item, marker) in enumerate(zip(expected_items, expected_markers)): + used_markers.add(marker) + expected_lines.append(f"{marker} [{index}] {item!r}{RESET}") + + actual_lines: list[str] = [] + for index, (item, marker) in enumerate(zip(actual_items, actual_markers)): + used_markers.add(marker) + actual_lines.append(f"{marker} [{index}] {item!r}{RESET}") + + newline = "\n" + expected_string = f"{BOLD}Expected items:{RESET}\n [\n{newline.join(expected_lines)}\n ]" + actual_string = f"{BOLD}Actually received items:{RESET}\n [\n{newline.join(actual_lines)}\n ]" + first_message = "Sequences are not equal" + + legend = f"{BOLD}Legend:{RESET}\n" + if CORRECT_MARKER in used_markers: + legend += f" {CORRECT_MARKER}{RESET}: item is correct as it was both expected and received\n" + if MISSING_MARKER in used_markers: + legend += ( + f" {MISSING_MARKER}{RESET}: item is expected but was not received\n" + ) + if UNEXPECTED_MARKER in used_markers: + legend += ( + f" {UNEXPECTED_MARKER}{RESET}: item was received but is not expected\n" + ) + if MOVED_MARKER in used_markers: + legend += f" {MOVED_MARKER}{RESET}: item is present in both sequences but at different positions\n" + + diff_message = ( + f"{first_message}\n{legend}\n{expected_string}\n\n{actual_string}" + ) + + extra_message = "" + if message is not None: + extra_message = f"\n{message}" self.fail(f"{diff_message}{extra_message}") @@ -413,6 +538,7 @@ class TestCase(unittest.TestCase): Specifically: - better errors for set inequality + - better errors for list inequality """ if first == second: return @@ -426,6 +552,19 @@ class TestCase(unittest.TestCase): exact=True, ) + if ( + isinstance(first, Sequence) + and isinstance(second, Sequence) + and all(_is_hashable(item) for item in first) + and all(_is_hashable(item) for item in second) + ): + self._fail_with_sequence_inequality( + first, + second, + # Any `str()`-able object is valid as a message. We frequently pass in raw JSON responses, for instance. + str(msg) if msg is not None else msg, + ) + # Fall back to the base implementation for other types # Since we know `first == second` does not hold, # we expect this must diverge by raising an exception. @@ -460,6 +599,32 @@ class TestCase(unittest.TestCase): self._fail_with_set_inequality(actual_items, expected_items, message, exact) + def assertEqualNormalisingWhitespace( + self, received: str, expected: str, msg: str | None = None + ) -> None: + """ + Fail the test if `received` and `expected` are not equal, + after having normalised all whitespace. + + By normalising whitespace, we mean that any run of whitespace + is replaced by a single ` `, on both sides. + The front and back of the strings are also trimmed. + + Whitespace characters considered are: `\n`, `\t`, ` `. + """ + PATTERN = r"[ \n\t]+" + REPLACEMENT = " " + normalised_received = re.sub(PATTERN, REPLACEMENT, received).strip() + normalised_expected = re.sub(PATTERN, REPLACEMENT, expected).strip() + + if normalised_received == normalised_expected: + return + + msg = "" if msg is None else msg + self.fail( + f"Expected strings to match, after normalising whitespace: {msg}\nReceived: {received}\nExpected: {expected}\nReceived (normalised): {normalised_received}\nExpected (normalised): {normalised_expected}" + ) + def DEBUG(target: TV) -> TV: """A decorator to set the .loglevel attribute to logging.DEBUG.