From 867780b5006100b1dc2dcb3083e8fcc30de3680e Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Tue, 30 Jun 2026 12:45:38 -0400 Subject: [PATCH 1/2] fix: ignore invalid batch write affinity candidates --- alternator/core/key_affinity.py | 82 +++++++------ tests/unit/test_key_affinity.py | 204 ++++++++++++++++++++++++++++++++ 2 files changed, 247 insertions(+), 39 deletions(-) diff --git a/alternator/core/key_affinity.py b/alternator/core/key_affinity.py index 6108733..091f565 100644 --- a/alternator/core/key_affinity.py +++ b/alternator/core/key_affinity.py @@ -31,6 +31,7 @@ class _BatchWriteRoutingTarget(NamedTuple): class _BatchWriteCandidate(NamedTuple): table_name: str attributes: dict[str, Any] + operation: str class AffinitySelector: @@ -254,22 +255,37 @@ def _iter_batch_write_candidates( if not isinstance(table_name, str) or not isinstance(writes, list): continue for write in writes: - if not isinstance(write, dict): - continue + candidate = _batch_write_candidate(table_name, write) + if candidate is not None: + candidates.append(candidate) - put_request = write.get("PutRequest") - if isinstance(put_request, dict): - item = put_request.get("Item") - if isinstance(item, dict): - candidates.append(_BatchWriteCandidate(table_name, item)) + return tuple(candidates) - delete_request = write.get("DeleteRequest") - if isinstance(delete_request, dict): - key = delete_request.get("Key") - if isinstance(key, dict): - candidates.append(_BatchWriteCandidate(table_name, key)) - return tuple(candidates) +def _batch_write_candidate( + table_name: str, + write: object, +) -> _BatchWriteCandidate | None: + if not isinstance(write, dict): + return None + + operations = [ + operation for operation in ("PutRequest", "DeleteRequest") if operation in write + ] + if len(operations) != 1: + return None + + operation = operations[0] + request = write[operation] + if not isinstance(request, dict): + return None + + attribute_field = "Item" if operation == "PutRequest" else "Key" + attributes = request.get(attribute_field) + if not isinstance(attributes, dict): + return None + + return _BatchWriteCandidate(table_name, attributes, operation) def _non_empty_string(value: object) -> bool: @@ -324,34 +340,22 @@ def _find_batch_write_routing_target( if not isinstance(table_name, str) or not isinstance(writes, list): continue for write in writes: - if not isinstance(write, dict): + candidate = _batch_write_candidate(table_name, write) + if candidate is None: continue - put_request = write.get("PutRequest") - if isinstance(put_request, dict): - item = put_request.get("Item") - if isinstance(item, dict): - target = _min_batch_write_target( - target, - _BatchWriteRoutingTarget( - table_name, - item, - _batch_write_sort_key(table_name, "PutRequest", item), - ), - ) - - delete_request = write.get("DeleteRequest") - if isinstance(delete_request, dict): - key = delete_request.get("Key") - if isinstance(key, dict): - target = _min_batch_write_target( - target, - _BatchWriteRoutingTarget( - table_name, - key, - _batch_write_sort_key(table_name, "DeleteRequest", key), - ), - ) + target = _min_batch_write_target( + target, + _BatchWriteRoutingTarget( + table_name, + candidate.attributes, + _batch_write_sort_key( + table_name, + candidate.operation, + candidate.attributes, + ), + ), + ) return target diff --git a/tests/unit/test_key_affinity.py b/tests/unit/test_key_affinity.py index fff7aa9..2d8868f 100644 --- a/tests/unit/test_key_affinity.py +++ b/tests/unit/test_key_affinity.py @@ -565,6 +565,141 @@ def test_batch_write_tied_votes_fall_back(self) -> None: is None ) + def test_batch_write_unknown_table_metadata_does_not_block_known_table( + self, + ) -> None: + """Unknown table metadata is skipped while another table can vote.""" + nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") + b1 = _pk_value_for_node(nodes, "b", "known-table") + params = { + "RequestItems": { + "unknown": [ + {"PutRequest": {"Item": {"pk": {"S": "ignored"}}}}, + ], + "orders": [ + {"DeleteRequest": {"Key": {"pk": {"S": b1}}}}, + ], + } + } + + assert ( + select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) + == "b" + ) + + def test_batch_write_unusable_candidates_are_skipped(self) -> None: + """Missing and unsupported keys are skipped while valid candidates vote.""" + nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") + c1 = _pk_value_for_node(nodes, "c", "usable") + params = { + "RequestItems": { + "orders": [ + {"PutRequest": {"Item": {"pk": {"BOOL": True}}}}, + {"DeleteRequest": {"Key": {"other": {"S": "missing-pk"}}}}, + {"PutRequest": {"Item": {"pk": {"S": c1}}}}, + ], + } + } + + assert ( + select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) + == "c" + ) + + def test_batch_write_non_key_attributes_do_not_change_winner(self) -> None: + """Non-key payload attributes do not affect batch affinity voting.""" + nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") + b1 = _pk_value_for_node(nodes, "b", "payload-b1") + b2 = _pk_value_for_node(nodes, "b", "payload-b2") + a1 = _pk_value_for_node(nodes, "a", "payload-a1") + params_a = { + "RequestItems": { + "orders": [ + { + "PutRequest": { + "Item": {"pk": {"S": b1}, "note": {"S": "before"}} + } + }, + { + "PutRequest": { + "Item": {"pk": {"S": b2}, "note": {"S": "before"}} + } + }, + { + "DeleteRequest": { + "Key": {"pk": {"S": a1}, "note": {"S": "before"}} + } + }, + ], + } + } + params_b = { + "RequestItems": { + "orders": [ + { + "DeleteRequest": { + "Key": {"pk": {"S": a1}, "note": {"S": "after"}} + } + }, + {"PutRequest": {"Item": {"pk": {"S": b2}, "note": {"S": "after"}}}}, + {"PutRequest": {"Item": {"pk": {"S": b1}, "note": {"S": "after"}}}}, + ], + } + } + + for params in (params_a, params_b): + assert ( + select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) + == "b" + ) + + def test_batch_write_invalid_union_write_request_is_ignored(self) -> None: + """A malformed write containing put and delete does not contribute votes.""" + nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") + a1 = _pk_value_for_node(nodes, "a", "valid-write") + b1 = _pk_value_for_node(nodes, "b", "invalid-put") + b2 = _pk_value_for_node(nodes, "b", "invalid-delete") + params = { + "RequestItems": { + "orders": [ + {"PutRequest": {"Item": {"pk": {"S": a1}}}}, + { + "PutRequest": {"Item": {"pk": {"S": b1}}}, + "DeleteRequest": {"Key": {"pk": {"S": b2}}}, + }, + ], + } + } + + assert ( + select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) + == "a" + ) + def test_batch_write_binary_pk_selects_stable_node(self) -> None: """Test binary partition-key values use stable hashing.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") @@ -661,6 +796,59 @@ def compute_affinity_node( assert first_url == "http://b:8000/" assert {second_url, third_url} == {"http://a:8000/", "http://c:8000/"} + def test_batch_write_preferred_node_first_and_remaining_nodes_preserved( + self, + ) -> None: + """Test BatchWriteItem affinity keeps retries after the preferred node.""" + config = Config(seed_hosts=["seed"], port=8000) + manager = MagicMock() + manager.nodes = NodeList(nodes=("a", "b", "c"), scope_name="cluster") + events = MagicMock() + params = { + "RequestItems": { + "orders": [ + {"PutRequest": {"Item": {"pk": {"S": "order-1"}}}}, + ], + } + } + + def compute_affinity_node( + operation_name: str, + request_params: dict[str, Any], + nodes: NodeList, + ) -> str | None: + assert operation_name == "BatchWriteItem" + assert request_params == params + assert nodes.nodes == ("a", "b", "c") + return "c" + + _register_alternator_handlers( + events, + manager, + config, + compute_affinity_node, + ) + handlers = { + call[0][1].__name__: call[0][1] for call in events.register.call_args_list + } + + request = MagicMock() + request.url = "http://seed:8000/" + request.headers = {"X-Amz-Target": "DynamoDB_20120810.BatchWriteItem"} + request.body = b'{"RequestItems":{"orders":[{"PutRequest":{"Item":{"pk":{"S":"order-1"}}}}]}}' + request._alternator_query_plan = None + + update_endpoint = handlers["update_endpoint"] + update_endpoint(request) + first_url = request.url + update_endpoint(request) + second_url = request.url + update_endpoint(request) + third_url = request.url + + assert first_url == "http://c:8000/" + assert {second_url, third_url} == {"http://a:8000/", "http://b:8000/"} + class TestExtractPartitionKey: """Tests for extract_partition_key function.""" @@ -904,6 +1092,22 @@ def test_extract_from_batch_get_shape(self) -> None: result = extract_partition_key(params, "user_id") assert result is None + def test_extract_from_invalid_batch_write_union_is_ignored(self) -> None: + """Batch write entries with both operations are not routing targets.""" + params = { + "RequestItems": { + "orders": [ + { + "PutRequest": {"Item": {"pk": {"S": "put"}}}, + "DeleteRequest": {"Key": {"pk": {"S": "delete"}}}, + }, + ], + } + } + + assert get_table_name(params) is None + assert extract_partition_key(params, "pk") is None + def test_key_not_found(self) -> None: """Test when partition key is not in params.""" params = {"Key": {"other_key": {"S": "value"}}} From 4a4274fe5a2efebbc06c0a4df392a5afc0a26039 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Tue, 30 Jun 2026 13:59:31 -0400 Subject: [PATCH 2/2] Align BatchWriteItem affinity query plan --- alternator/async_client.py | 5 +- alternator/client.py | 5 +- alternator/core/handlers.py | 24 ++++- alternator/core/key_affinity.py | 16 +-- tests/unit/test_key_affinity.py | 185 ++++++++++++++------------------ 5 files changed, 113 insertions(+), 122 deletions(-) diff --git a/alternator/async_client.py b/alternator/async_client.py index dc1541c..fd7e380 100644 --- a/alternator/async_client.py +++ b/alternator/async_client.py @@ -25,6 +25,7 @@ from alternator.core.auth import apply_auth from alternator.core.handlers import _register_alternator_handlers from alternator.core.key_affinity import ( + AffinityTarget, select_affinity_node, ) from alternator.core.live_nodes import AsyncLiveNodesManager, NodeList @@ -203,7 +204,7 @@ def preload(self, table_pk_map: dict[str, str]) -> None: def _create_async_affinity_node_computer( config: Config, pk_cache: AsyncPartitionKeyCache | None, -) -> Callable[[str, dict[str, Any], NodeList], str | None] | None: +) -> Callable[[str, dict[str, Any], NodeList], AffinityTarget | None] | None: """ Create a function that selects the preferred key-affinity node. @@ -240,7 +241,7 @@ def compute_affinity_node( operation_name: str, params: dict[str, Any], nodes: NodeList, - ) -> str | None: + ) -> AffinityTarget | None: """Select the preferred key-affinity node for this request.""" return select_affinity_node( mode=affinity_mode.name, diff --git a/alternator/client.py b/alternator/client.py index 385581d..0c2c3c4 100644 --- a/alternator/client.py +++ b/alternator/client.py @@ -20,6 +20,7 @@ from alternator.core.auth import apply_auth from alternator.core.handlers import _register_alternator_handlers from alternator.core.key_affinity import ( + AffinityTarget, PartitionKeyCache, select_affinity_node, ) @@ -167,7 +168,7 @@ def _create_boto_config(config: Config, *, auth_enabled: bool) -> BotoConfig: def _create_affinity_node_computer( config: Config, client: DynamoDBClient, -) -> Callable[[str, dict[str, Any], NodeList], str | None] | None: +) -> Callable[[str, dict[str, Any], NodeList], AffinityTarget | None] | None: """ Create a function that selects the preferred key-affinity node. @@ -198,7 +199,7 @@ def compute_affinity_node( operation_name: str, params: dict[str, Any], nodes: NodeList, - ) -> str | None: + ) -> AffinityTarget | None: """Select the preferred key-affinity node for this request.""" return select_affinity_node( mode=affinity_mode.name, diff --git a/alternator/core/handlers.py b/alternator/core/handlers.py index 75b9a8b..1955551 100644 --- a/alternator/core/handlers.py +++ b/alternator/core/handlers.py @@ -22,6 +22,7 @@ create_header_filter_handler, create_user_agent_header_handler, ) +from alternator.core.key_affinity import AffinityTarget from alternator.core.query_plan import LazyQueryPlan from alternator.core.request import extract_operation_name, extract_request_params from alternator.exceptions import NoNodesAvailableError @@ -49,7 +50,9 @@ def _register_alternator_handlers( events: BaseEventHooks, manager: _HasNodes, config: Config, - compute_affinity_node: Callable[[str, DynamoDBParams, NodeList], str | None] + compute_affinity_node: Callable[ + [str, DynamoDBParams, NodeList], AffinityTarget | None + ] | None = None, *, auth_enabled: bool = False, @@ -78,10 +81,21 @@ def _register_alternator_handlers( def create_query_plan( nodes: NodeList, - preferred_node: str | None, + preferred_node: AffinityTarget | None, ) -> Iterator[str]: """Create a URI iterator for a single request.""" node_addresses = nodes.nodes + if isinstance(preferred_node, tuple): + emitted: set[str] = set() + for node in preferred_node: + if node in node_addresses and node not in emitted: + emitted.add(node) + yield f"{scheme}://{node}:{port}" + for node in node_addresses: + if node not in emitted: + yield f"{scheme}://{node}:{port}" + return + if preferred_node is not None and preferred_node in node_addresses: yield f"{scheme}://{preferred_node}:{port}" remaining_nodes = tuple( @@ -146,7 +160,7 @@ def update_endpoint( def _create_request_query_plan( request: AWSRequest | AWSPreparedRequest, - preferred_node: str | None | object = _PREFERRED_NODE_UNSET, + preferred_node: AffinityTarget | None | object = _PREFERRED_NODE_UNSET, ) -> Iterator[str]: nodes = manager.nodes if not nodes: @@ -155,7 +169,7 @@ def _create_request_query_plan( scope_name=scope_name, ) - selected_preferred_node: str | None + selected_preferred_node: AffinityTarget | None if preferred_node is _PREFERRED_NODE_UNSET: selected_preferred_node = None if compute_affinity_node is not None: @@ -170,7 +184,7 @@ def _create_request_query_plan( nodes, ) else: - selected_preferred_node = cast("str | None", preferred_node) + selected_preferred_node = cast("AffinityTarget | None", preferred_node) return create_query_plan(nodes, selected_preferred_node) def _store_query_plan( diff --git a/alternator/core/key_affinity.py b/alternator/core/key_affinity.py index 091f565..931aa30 100644 --- a/alternator/core/key_affinity.py +++ b/alternator/core/key_affinity.py @@ -20,6 +20,7 @@ from alternator.core.live_nodes import NodeList logger = logging.getLogger("alternator") +AffinityTarget = str | tuple[str, ...] class _BatchWriteRoutingTarget(NamedTuple): @@ -123,7 +124,7 @@ def select_affinity_node( params: dict[str, Any], nodes: NodeList, get_pk_name: Callable[[str], str | None], -) -> str | None: +) -> AffinityTarget | None: """Select the preferred affinity node for a request, or None for fallback.""" if not should_use_affinity(mode, operation_name, params): return None @@ -204,7 +205,7 @@ def _select_batch_write_affinity_node( params: dict[str, Any], nodes: NodeList, get_pk_name: Callable[[str], str | None], -) -> str | None: +) -> tuple[str, ...] | None: votes: Counter[str] = Counter() selector = AffinitySelector() @@ -234,11 +235,12 @@ def _select_batch_write_affinity_node( if not votes: return None - top_count = max(votes.values()) - winners = [node for node, count in votes.items() if count == top_count] - if len(winners) != 1: - return None - return winners[0] + return tuple( + sorted( + votes, + key=lambda node: (-votes[node], node), + ) + ) def _iter_batch_write_candidates( diff --git a/tests/unit/test_key_affinity.py b/tests/unit/test_key_affinity.py index 2d8868f..e8107b1 100644 --- a/tests/unit/test_key_affinity.py +++ b/tests/unit/test_key_affinity.py @@ -367,16 +367,13 @@ def test_batch_write_single_put_selects_node(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "a" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("a",) def test_batch_write_single_delete_selects_node(self) -> None: """Test BatchWriteItem with a single DeleteRequest selects its node.""" @@ -388,16 +385,13 @@ def test_batch_write_single_delete_selects_node(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "c" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("c",) def test_batch_write_mixed_put_delete_unique_winner(self) -> None: """Test BatchWriteItem votes for the unique preferred node.""" @@ -415,16 +409,13 @@ def test_batch_write_mixed_put_delete_unique_winner(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "b" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("b", "c") def test_batch_write_multi_table_reversed_order_same_winner(self) -> None: """Test batch voting is independent of table and request order.""" @@ -446,16 +437,13 @@ def test_batch_write_multi_table_reversed_order_same_winner(self) -> None: } for params in (params_a, params_b): - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk", "sessions": "pk"}.get, - ) - == "b" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk", "sessions": "pk"}.get, + ) == ("b", "a") def test_batch_write_missing_pk_metadata_falls_back(self) -> None: """Test missing partition-key metadata produces no preferred node.""" @@ -540,8 +528,8 @@ def test_batch_write_no_nodes_falls_back(self) -> None: is None ) - def test_batch_write_tied_votes_fall_back(self) -> None: - """Test tied preferred-node votes produce no preferred node.""" + def test_batch_write_tied_votes_use_node_address_tie_break(self) -> None: + """Test tied preferred-node votes use node address order.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") a1 = _pk_value_for_node(nodes, "a", "tie-a") b1 = _pk_value_for_node(nodes, "b", "tie-b") @@ -554,16 +542,13 @@ def test_batch_write_tied_votes_fall_back(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - is None - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("a", "b") def test_batch_write_unknown_table_metadata_does_not_block_known_table( self, @@ -582,16 +567,13 @@ def test_batch_write_unknown_table_metadata_does_not_block_known_table( } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "b" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("b",) def test_batch_write_unusable_candidates_are_skipped(self) -> None: """Missing and unsupported keys are skipped while valid candidates vote.""" @@ -607,16 +589,13 @@ def test_batch_write_unusable_candidates_are_skipped(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "c" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("c",) def test_batch_write_non_key_attributes_do_not_change_winner(self) -> None: """Non-key payload attributes do not affect batch affinity voting.""" @@ -660,16 +639,13 @@ def test_batch_write_non_key_attributes_do_not_change_winner(self) -> None: } for params in (params_a, params_b): - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "b" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("b", "a") def test_batch_write_invalid_union_write_request_is_ignored(self) -> None: """A malformed write containing put and delete does not contribute votes.""" @@ -689,16 +665,13 @@ def test_batch_write_invalid_union_write_request_is_ignored(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == "a" - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == ("a",) def test_batch_write_binary_pk_selects_stable_node(self) -> None: """Test binary partition-key values use stable hashing.""" @@ -716,16 +689,13 @@ def test_batch_write_binary_pk_selects_stable_node(self) -> None: } } - assert ( - select_affinity_node( - mode="ANY_WRITE", - operation_name="BatchWriteItem", - params=params, - nodes=nodes, - get_pk_name={"orders": "pk"}.get, - ) - == expected - ) + assert select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) == (expected,) def test_batch_write_selection_does_not_mutate_params(self) -> None: """Test BatchWriteItem affinity selection leaves request params unchanged.""" @@ -816,11 +786,11 @@ def compute_affinity_node( operation_name: str, request_params: dict[str, Any], nodes: NodeList, - ) -> str | None: + ) -> tuple[str, ...] | None: assert operation_name == "BatchWriteItem" assert request_params == params assert nodes.nodes == ("a", "b", "c") - return "c" + return ("c", "a") _register_alternator_handlers( events, @@ -846,8 +816,11 @@ def compute_affinity_node( update_endpoint(request) third_url = request.url - assert first_url == "http://c:8000/" - assert {second_url, third_url} == {"http://a:8000/", "http://b:8000/"} + assert (first_url, second_url, third_url) == ( + "http://c:8000/", + "http://a:8000/", + "http://b:8000/", + ) class TestExtractPartitionKey: