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.