import pathlib from time import sleep from typing import Any, Literal from .assertions import assert_hw_measurements_equal_many from .fixtures import upsert_random_points, create_collection, get_telemetry_hw_info from .utils import * COLLECTION_NAME = "test_collection" def test_resharding_state_transitions(tmp_path: pathlib.Path): """ Tests allowed state transitions during resharding """ # Bootstrap resharding cluster peer_uris, _ = bootstrap_resharding(tmp_path) # Check that we can't (re)apply unexpected resharding state transitions try_requests(peer_uris[0], 400, [ start_resharding, lambda peer_uri: start_resharding(peer_uri, direction="down"), commit_write_hashring, finish_resharding, ]) # Commit read hashring resp = commit_read_hashring(peer_uris[0]) assert_http_ok(resp) # Check that we can't (re)apply unexpected resharding state transitions or abort resharding try_requests(peer_uris[0], 400, [ start_resharding, lambda peer_uri: start_resharding(peer_uri, direction="down"), finish_resharding, abort_resharding, ]) # Commit write hashring resp = commit_write_hashring(peer_uris[0]) assert_http_ok(resp) # Check that we can't (re)apply unexpected resharding state transitions or abort resharding try_requests(peer_uris[0], 400, [ start_resharding, lambda peer_uri: start_resharding(peer_uri, direction="down"), abort_resharding, ]) # Finish resharding resp = finish_resharding(peer_uris[0]) assert_http_ok(resp) # Wait for resharding to finish wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) def test_resharding_abort(tmp_path: pathlib.Path): """ Tests that resharding can be aborted """ # Bootstrap resharding cluster peer_uris, _ = bootstrap_resharding(tmp_path) # Abort resharding resp = abort_resharding(peer_uris[0]) assert_http_ok(resp) # Wait for resharding to abort wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) def test_resharding_abort_with_replicas(tmp_path: pathlib.Path): """ Tests that resharding can be aborted after replication Covers bug . """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding( tmp_path, shard_number=1, replication_factor=3, peers=3, upsert_points=1000, ) # Migrate resharding points from shard 0 into 1 migrate_points(peer_uris[0], peer_ids[0], 0, peer_ids[0], 1, "up") activate_replica(peer_uris[0], peer_ids[0], 1) # Replicate new shard twice to match replication factor replicate_shard(peer_uris[0], peer_ids[0], peer_ids[1], 1) replicate_shard(peer_uris[0], peer_ids[0], peer_ids[2], 1) # Assert that resharding is still in progress info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) assert "resharding_operations" in info and len(info["resharding_operations"]) == 1 # Abort resharding resp = abort_resharding(peer_uris[0]) assert_http_ok(resp) # Wait for resharding to abort wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) def test_resharding_abort_on_delete_collection(tmp_path: pathlib.Path): """ Testa that resharding is automatically aborted, when collection is deleted """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path, peer_idx=-1) # Remove collection resp = requests.delete(f"{peer_uris[0]}/collections/{COLLECTION_NAME}") assert_http_ok(resp) # TODO: Wait for/check... *what*? 🤔 def test_resharding_abort_on_delete_shard_key(tmp_path: pathlib.Path): """ Tests that resharding is automatically aborted, when custom shard key is deleted """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding( tmp_path, shard_keys=["custom_shard_key_1", "custom_shard_key_2"], shard_key="custom_shard_key_2", ) # Delete shard key resp = requests.post(f"{peer_uris[0]}/collections/{COLLECTION_NAME}/shards/delete", json={ "shard_key": "custom_shard_key_2", }) assert_http_ok(resp) # Wait for resharding to abort wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) def test_resharding_abort_on_remove_peer(tmp_path: pathlib.Path): """ Tests that resharding is automatically aborted, when we force-remove resharding peer """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path, peer_idx=-1) # Remove target peer resp = requests.delete(f"{peer_uris[0]}/cluster/peer/{peer_ids[-1]}?force=true") assert_http_ok(resp) # Wait for resharding to abort wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) def test_resharding_try_remove_target_shard(tmp_path: pathlib.Path): """ Tests that new shard can't be removed during resharding (until it has been replicated) """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path) # Try to remove target shard resp = requests.post(f"{peer_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ "drop_replica": { "peer_id": peer_ids[0], "shard_id": 3, } }) assert_http(resp, 400) @pytest.mark.parametrize("direction, peers", [("up", 3), ("down", 3)]) def test_resharding_forward(tmp_path: pathlib.Path, direction: Literal["up", "down"], peers: int): """ Tests that updates are forwarded to target replicas """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path, direction=direction, peers=peers) # Upsert points to collection upsert_random_points(peer_uris[0], 1000, collection_name=COLLECTION_NAME) # Get collection cluster info info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) # Select target replica target_shard_id = peers if direction == "up" else peers - 1 target_peer_id, target_peer_uri = find_replica(target_shard_id, info, peer_uris, peer_ids) # Create a list to collect replica URIs selected during test replica_uris = [] for shard_id in range(target_shard_id): # Find replicas of selected shard peer_id, peer_uri = find_replica(shard_id, info, peer_uris, peer_ids) # Assert that all points were correctly forwarded assert_resharding_points(peer_uri, shard_id, target_peer_uri, target_shard_id) # Append peer URI to the list of replica URIs replica_uris.append(peer_uri) # Append target replica to the list of replica URIs replica_uris.append(target_peer_uri) # Assert total count of migrated points assert_resharding_points_count(replica_uris) @pytest.mark.parametrize("direction, peers", [("up", 3), ("down", 3)]) def test_resharding_transfer(tmp_path: pathlib.Path, direction: Literal["up", "down"], peers: int): """ Tests that resharding transfers migrate points to target replicas """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path, upsert_points=1000, direction=direction, peers=peers) # Get collection cluster info info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) # Select target replica target_shard_id = peers if direction == "up" else peers - 1 target_peer_id, target_peer_uri = find_replica(target_shard_id, info, peer_uris, peer_ids) # Create a list to collect replica URIs selected during test replica_uris = [] for shard_id in range(target_shard_id): # Find replica of selected shard peer_id, peer_uri = find_replica(shard_id, info, peer_uris, peer_ids) # Collect all nodes hardware measurements before transferring hw = [get_telemetry_hw_info(uri, COLLECTION_NAME) for uri in peer_uris] # Migrate resharding points migrate_points(peer_uris[0], peer_id, shard_id, target_peer_id, target_shard_id, direction) # Assert that no hardware measurements have been measured for the transfer new_hw = [get_telemetry_hw_info(uri, COLLECTION_NAME) for uri in peer_uris] assert_hw_measurements_equal_many(hw, new_hw) # Assert that all points were correctly migrated assert_resharding_points(peer_uri, shard_id, target_peer_uri, target_shard_id) # Append peer URI to the list of replica URIs replica_uris.append(peer_uri) # Append target replica to the list of replica URIs replica_uris.append(target_peer_uri) # Assert total count of migrated points assert_resharding_points_count(replica_uris) def test_resharding_clean_update_queue_data_race(tmp_path: pathlib.Path): """ Tests data race between resharding clean operation and the update queue. During resharding up, after both hash rings are committed, the cleanup operation removes points from existing shards that, under the new hash ring, now belong on the new shard. However, when the update queue is in use (wait=false), there may be pending operations waiting in the queue when cleanup begins. The cleanup task scrolls segment storage directly - it does not see anything that is still in the queue. Its delete operation then targets only what it scrolled. Anything in the queue that lands on the shard *after* cleanup has scrolled is missed and survives the cleanup as an invalid point. To deterministically trigger the race we: - Queue the upserts during the MigratingPoints stage. At this stage the sender (old) shards still receive normal upserts (the `update_only_existing` split kicks in only at ReadHashRingCommitted and later), so points routed via the resharding ring genuinely land on the old shards. - Use a staging delay on every peer to block each shard's worker. The delay does not forward over gRPC (see remote_shard.rs), so it must be sent to each peer individually. - Use a long delay so that all cleanup calls are issued while every worker is still blocked - ensuring the scroll runs before the queued upserts apply. """ assert_project_root() # Bootstrap cluster: 3 peers, 3 shards, 1 replica each, with some initial points peer_uris, peer_ids = bootstrap_cluster( tmp_path, shard_number=3, replication_factor=1, peers=3, upsert_points=300, ) # Need staging feature for the delay debug operation skip_if_no_feature(peer_uris[0], "staging") # Start resharding up - this adds shard 3, stage = MigratingPoints resp = start_resharding(peer_uris[0]) assert_http_ok(resp) wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 1) info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) target_shard_id = 3 target_peer_id, _ = find_replica(target_shard_id, info, peer_uris, peer_ids) # Migrate the existing points from each old shard to the new shard for shard_id in range(target_shard_id): peer_id, _ = find_replica(shard_id, info, peer_uris, peer_ids) migrate_points(peer_uris[0], peer_id, shard_id, target_peer_id, target_shard_id, "up") # Activate the new shard's replica - stage is still MigratingPoints activate_replica(peer_uris[0], target_peer_id, target_shard_id) # Inject a delay on EVERY peer. The staging delay is broadcast to all # shards but skipped by gRPC for remote shards (remote_shard.rs:878), # so we must hit each peer directly to actually block every shard's # update worker. A long delay keeps every queue blocked across the # entire window in which cleanups are issued. DELAY_SEC = 15.0 for uri in peer_uris: resp = requests.post( f"{uri}/collections/{COLLECTION_NAME}/debug", json={"delay": {"duration_sec": DELAY_SEC}}, ) assert_http_ok(resp) # Insert NEW points with wait=false WHILE STAGE IS STILL MigratingPoints. # In this stage, sender shards (the existing ones) are NOT subject to # the update_only_existing split - they receive normal upserts via the # resharding ring. Points routed here that, under the new ring, belong # to shard 3 will genuinely land on the old shards once the delay # clears - exactly the invalid points that cleanup must remove. upsert_random_points( peer_uris[0], 200, offset=10_000, wait="false", batch_size=10, ) # Commit read hash ring (stage = ReadHashRingCommitted, queues unaffected) resp = commit_read_hashring(peer_uris[0]) assert_http_ok(resp) # Commit write hash ring (ring becomes Single(new), queues unaffected) resp = commit_write_hashring(peer_uris[0]) assert_http_ok(resp) # Trigger cleanup on each existing shard with wait=true. The cleanup # task scrolls segment storage - which does not see queued operations - # then queues a delete for the IDs it scrolled. Because every worker is # still blocked by the delay, the cleanup's delete is queued AFTER the # pending upserts. When the delay finally clears, the worker applies the # upserts first (inserting the invalid points into segments), then the # delete (which only targets the originally-scrolled IDs, missing the # just-inserted invalid points). for shard_id in range(target_shard_id): _, peer_uri = find_replica(shard_id, info, peer_uris, peer_ids) resp = requests.post( f"{peer_uri}/collections/{COLLECTION_NAME}/shards/{shard_id}/cleanup", params={"wait": "true", "timeout": "60"}, ) assert_http_ok(resp) assert resp.json()["result"]["status"] == "completed" # Synchronize: a wait=true upsert ensures the queue is fully drained # everywhere before we measure. upsert_random_points(peer_uris[0], 1, offset=99_999, wait="true") # Verify: no existing shard should hold any point that, under the new # hash ring, belongs on the new shard. Any non-zero count exposes the # data race - an insert that landed on an old shard after cleanup ran. for shard_id in range(target_shard_id): _, peer_uri = find_replica(shard_id, info, peer_uris, peer_ids) invalid_count = count_local_points( peer_uri, shard_id, filter_shard_id=target_shard_id, exact=True, ) assert invalid_count == 0, ( f"shard {shard_id} on {peer_uri} contains {invalid_count} points " f"that belong to new shard {target_shard_id} after cleanup; this " f"indicates a data race between cleanup and the update queue" ) # Tidy up: finish the resharding operation resp = finish_resharding(peer_uris[0]) assert_http_ok(resp) wait_for_collection_resharding_operations_count(peer_uris[0], COLLECTION_NAME, 0) @pytest.mark.parametrize("peers", [(3)]) def test_resharding_down_abort_cleanup(tmp_path: pathlib.Path, peers: int): """ Tests that migrated points are cleaned up from target replicas, when scale-down resharding is aborted """ # Bootstrap resharding cluster peer_uris, peer_ids = bootstrap_resharding(tmp_path, upsert_points=500, direction="down", peers=peers) # Upsert points to collection upsert_random_points(peer_uris[0], 500, collection_name=COLLECTION_NAME) # Get collection cluster info info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) # Select target replica target_shard_id = peers - 1 target_peer_id, _ = find_replica(target_shard_id, info, peer_uris, peer_ids) # Create a list to collect replica URIs selected during test replica_uris = [] for shard_id in range(target_shard_id): # Find replica of selected shard peer_id, peer_uri = find_replica(shard_id, info, peer_uris, peer_ids) # Migrate resharding points migrate_points(peer_uris[0], peer_id, shard_id, target_peer_id, target_shard_id, "down") # Assert that some points were forwarded and/or migrated to selected replica resharding_points_count = count_local_points(peer_uri, shard_id, target_shard_id, exact=True) assert resharding_points_count > 0 # Append peer URI to the list of replica URIs replica_uris.append(peer_uri) # Abort resharding resp = abort_resharding(peer_uris[0]) assert_http_ok(resp) # Wait for resharding to abort on all peers for peer_uri in replica_uris: wait_for_collection_resharding_operations_count(peer_uri, COLLECTION_NAME, 0) # Assert that all replicas are in `Active` state info = get_collection_cluster_info(peer_uris[0], COLLECTION_NAME) for replica in all_replicas(info): assert replica["state"] == "Active" # Assert that forwarded and/or migrated points were deleted from non-target replicas for shard_id, peer_uri in enumerate(replica_uris): resharding_points_count = count_local_points(peer_uri, shard_id, target_shard_id, exact=True) assert resharding_points_count == 0 def bootstrap_resharding( tmp_path: pathlib.Path, collection: str = COLLECTION_NAME, peer_idx: int | None = None, **kwargs, ): # Bootstrap cluster peer_uris, peer_ids = bootstrap_cluster(tmp_path, collection, **kwargs) wait_for_all_peers_versions(peer_uris) # Select target peer peer_id = None if peer_idx: try: peer_id = peer_ids[peer_idx] finally: pass # Start resharding resp = start_resharding(peer_uris[0], collection, peer_id=peer_id, **kwargs) assert_http_ok(resp) # Wait for resharding to start wait_for_collection_resharding_operations_count(peer_uris[0], collection, 1) return (peer_uris, peer_ids) def bootstrap_cluster( tmp_path: pathlib.Path, collection: str = COLLECTION_NAME, shard_number: int = 3, replication_factor: int = 2, shard_keys: list[str] | str | None = None, upsert_points: int = 0, peers: int = 3, **kwargs, ) -> tuple[list[str], list[int]]: assert_project_root() # Prevent optimizers messing with point counts env = { "QDRANT__STORAGE__OPTIMIZERS__INDEXING_THRESHOLD_KB": "0", } # Start cluster peer_uris, _, _ = start_cluster(tmp_path, peers, extra_env=env) # Collect peer IDs peer_ids = [] for peer_uri in peer_uris: peer_ids.append(get_cluster_info(peer_uri)["peer_id"]) # Create collection create_collection( peer_uris[0], collection, shard_number, replication_factor, sharding_method="auto" if shard_keys is None else "custom", ) wait_collection_exists_and_active_on_all_peers(collection, peer_uris) # Create custom shard keys (if required), and upload points to collection if type(shard_keys) is not list: shard_keys: list[str | None] = [shard_keys] for shard_key in shard_keys: # Create custom shard key (if required) if shard_key is not None: resp = requests.put(f"{peer_uris[0]}/collections/{collection}/shards", json={ "shard_key": shard_key, "shards_number": shard_number, "replication_factor": replication_factor, }) assert_http_ok(resp) # Upsert points to collection if upsert_points > 0: upsert_random_points( peer_uris[0], upsert_points, collection_name=collection, shard_key=shard_key, ) return (peer_uris, peer_ids) def start_resharding( peer_uri: str, collection: str = COLLECTION_NAME, direction: Literal["up", "down"] = "up", peer_id: int | None = None, shard_key: str | None = None, **kwargs, ): return requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "start_resharding": { "direction": direction, "peer_id": peer_id, "shard_key": shard_key, } }) def commit_read_hashring(peer_uri: str, collection: str = COLLECTION_NAME): return requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "commit_read_hash_ring": {} }) def commit_write_hashring(peer_uri: str, collection: str = COLLECTION_NAME): return requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "commit_write_hash_ring": {} }) def finish_resharding(peer_uri: str, collection: str = COLLECTION_NAME): return requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "finish_resharding": {} }) def abort_resharding(peer_uri: str, collection: str = COLLECTION_NAME): return requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "abort_resharding": {} }) def get_local_points( peer_uri: str, shard_id: int, point_ids: list[int], collection: str = COLLECTION_NAME, ): resp = requests.post(f"{peer_uri}/collections/{collection}/shards/{shard_id}/points", json = { "ids": point_ids, "with_vector": True, "with_payload": True, }) assert_http_ok(resp) return resp.json()['result'] def scroll_local_points( peer_uri: str, shard_id: int, filter_shard_id: int | None = None, offset: int | None = None, limit: int | None = None, collection: str = COLLECTION_NAME, ): resp = requests.post(f"{peer_uri}/collections/{collection}/shards/{shard_id}/points/scroll", json={ "limit": limit, "offset": offset, "hash_ring_filter": None if filter_shard_id is None else { "expected_shard_id": filter_shard_id, }, "with_vector": True, "with_payload": True, }) assert_http_ok(resp) return resp.json()['result'] def count_local_points( peer_uri: str, shard_id: int, filter_shard_id: int | None = None, exact: bool = True, collection: str = COLLECTION_NAME, ) -> int: resp = requests.post(f"{peer_uri}/collections/{collection}/shards/{shard_id}/points/count", json={ "exact": exact, "hash_ring_filter": None if filter_shard_id is None else { "expected_shard_id": filter_shard_id, }, }) assert_http_ok(resp) return resp.json()['result']['count'] def migrate_points( peer_uri: str, peer_id: int, shard_id: int, target_peer_id: int, target_shard_id: int, direction: Literal["up", "down"], collection: str = COLLECTION_NAME, info: dict[str, Any] | None = None, ) -> tuple[int, int]: """ Migrates resharding points to or from `shard_id`, depending on resharding `direction`: - if `direction` is "up", migrates points from `shard_id` to target shard - if `direction` is "down", migrates points from target shard to `shard_id` """ # Select replicas for resharding transfer if direction == "up": from_shard_id = shard_id from_peer_id = peer_id to_shard_id = target_shard_id to_peer_id = target_peer_id else: from_shard_id = target_shard_id from_peer_id = target_peer_id to_shard_id = shard_id to_peer_id = peer_id # Start resharding transfer resp = requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "replicate_shard": { "from_peer_id": from_peer_id, "to_peer_id": to_peer_id, "shard_id": from_shard_id, "to_shard_id": to_shard_id, "method": "resharding_stream_records", } }) assert_http_ok(resp) # Wait for resharding transfer to start sleep(1) # Wait for resharding transfer to finish or abort wait_for_collection_shard_transfers_count(peer_uri, collection, 0) # Assert that resharding transfer finished successfully info = get_collection_cluster_info(peer_uri, collection) # Assert that resharding is still in progress assert "resharding_operations" in info and len(info["resharding_operations"]) > 0 # Assert that replica `to_shard_id`@`to_peer_id` is in `Resharding` state migration_successful = False for replica in all_replicas(info): if replica["shard_id"] == to_shard_id and replica["peer_id"] == to_peer_id and replica["state"] in ("Resharding", "ReshardingScaleDown"): migration_successful = True break assert migration_successful # Return replicas used for resharding transfer return (peer_id, target_peer_id) def activate_replica( peer_uri: str, peer_id: int, shard_id: int, collection: str = COLLECTION_NAME, ): """ Activate a resharding replica that as been migrated into. """ # Activate migrated replica resp = requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "finish_migrating_points": { "peer_id": peer_id, "shard_id": shard_id, } }) assert_http_ok(resp) # Wait for replica activation sleep(1) # Assert that resharding transfer finished successfully info = get_collection_cluster_info(peer_uri, collection) # Assert that resharding is still in progress assert "resharding_operations" in info and len(info["resharding_operations"]) > 0 # Assert that replica `to_shard_id`@`to_peer_id` is in `Active` state migration_successful = False for replica in all_replicas(info): if replica["shard_id"] == shard_id and replica["peer_id"] == peer_id and replica["state"] in ("Active"): migration_successful = True break assert migration_successful def replicate_shard( peer_uri: str, from_peer_id: int, to_peer_id: int, shard_id: int, collection: str = COLLECTION_NAME, ): """ Replicate a shard from `from_peer_id` to `to_peer_id` """ # Start resharding transfer resp = requests.post(f"{peer_uri}/collections/{collection}/cluster", json={ "replicate_shard": { "from_peer_id": from_peer_id, "to_peer_id": to_peer_id, "shard_id": shard_id, "method": "stream_records", } }) assert_http_ok(resp) # Wait for resharding transfer to start sleep(1) # Wait for resharding transfer to finish or abort wait_for_collection_shard_transfers_count(peer_uri, collection, 0) def assert_resharding_points(peer_uri: str, shard_id: int, target_peer_uri: str, target_shard_id: int): """ Asserts that target replica contains all resharding points from selected replica and these points are exactly the same between replicas """ offset = 0 while offset is not None: # Scroll resharding points in selected replica resp = scroll_local_points(peer_uri, shard_id, target_shard_id, offset, 1000) # Fetch the same points from target replica target_resp = get_local_points(target_peer_uri, target_shard_id, [point['id'] for point in resp['points']]) # Assert that target replica contains all resharding points from selected replica # and these points are exactly the same between replicas assert target_resp == resp['points'] offset = resp['next_page_offset'] def assert_resharding_points_count(replica_uris: list[str]): """ Asserts that target replica points count matches total resharding points count in all other replicas """ # Select target replica target_shard_id = len(replica_uris) - 1 target_shard_uri = replica_uris[-1] # Get points count in target replica target_points_count = count_local_points(target_shard_uri, target_shard_id, exact=True) # Calculate total resharding points count in all other replicas total_resharding_points_count = 0 for shard_id, shard_uri in enumerate(replica_uris[:-1]): total_resharding_points_count += count_local_points(shard_uri, shard_id, target_shard_id, exact=True) # Assert target replica points count matches total resharding points count assert target_points_count == total_resharding_points_count def find_replica(shard_id: int, info: dict[str, Any], peer_uris: list[str], peer_ids: list[int]) -> tuple[int, str]: for replica in all_replicas(info): if replica["shard_id"] == shard_id: peer_id = replica["peer_id"] peer_uri = peer_uris[peer_ids.index(peer_id)] return (peer_id, peer_uri) raise Exception(f"replica of shard {shard_id} not found: {info}") def all_replicas(info: dict[str, Any]): for local in info["local_shards"]: local["peer_id"] = info["peer_id"] yield local for remote in info["remote_shards"]: yield remote def try_requests( peer_uri: str, expected_status: int, reqs: list[Callable[[str], requests.Response]], ): """ Execute multiple requests and asserts that all requests return expected status code """ for req in reqs: resp = req(peer_uri) assert_http(resp, expected_status) def assert_http(resp: requests.Response, expected_status: int): assert resp.status_code == expected_status, ( f"`{resp.url}` " f"returned an unexpected status code (expected {expected_status}, received {resp.status_code}):\n" f"{resp.json()}" )