Wait for all shards to be active in when testing collection recovery (q…

…drant#1508)
godhelpjun · Mar 15, 2023 · 1e1a2d9 · 1e1a2d9
1 parent c2316db
commit 1e1a2d9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/lib/collection/src/shards/replica_set.rs b/lib/collection/src/shards/replica_set.rs
@@ -55,11 +55,11 @@ const REPLICA_STATE_FILE: &str = "replica_state.json";
 //    │    Collection Created
 //    │
 //    ▼
-//  ┌─────────────┐
-//  │             │
-//  │ Initilizing │
-//  │             │
-//  └──────┬──────┘
+//  ┌──────────────┐
+//  │              │
+//  │ Initializing │
+//  │              │
+//  └──────┬───────┘
 //         │  Report created    ┌───────────┐
 //         └────────────────────►           │
 //             Activate         │ Consensus │

diff --git a/tests/consensus_tests/test_collection_recovery.py b/tests/consensus_tests/test_collection_recovery.py
@@ -1,7 +1,6 @@
 import pathlib
 import shutil
 
-import requests
 from .fixtures import create_collection, upsert_random_points, random_vector, search
 from .utils import *
 from .assertions import assert_http_ok
@@ -42,7 +41,7 @@ def test_collection_recovery(tmp_path: pathlib.Path):
     # Recover Raft state
     recover_raft_state(peer_url)
 
-    # Wait for the Raft state to be recovered
+    # Wait for all shards to be active
     wait_for(all_collection_shards_are_active, peer_url, COLLECTION_NAME)
 
     # Check, that the collection is not empty on recovered node
@@ -52,18 +51,16 @@ def test_collection_recovery(tmp_path: pathlib.Path):
         assert shard["points_count"] > 0
 
 
-def get_collection_cluser_info(peer_url, collection_name):
-    r = request.get(f"{peer_url}/collections/{collection_name}/cluster")
-    return request_result(r)
-
 def recover_raft_state(peer_url):
     r = requests.post(f"{peer_url}/cluster/recover")
     return request_result(r)
 
+
 def request_result(resp):
     assert_http_ok(resp)
     return resp.json()["result"]
 
+
 def collection_exists(peer_url, collection_name):
     try:
         get_collection_cluster_info(peer_url, collection_name)
@@ -72,15 +69,17 @@ def collection_exists(peer_url, collection_name):
 
     return True
 
+
 def all_collection_shards_are_active(peer_url, collection_name):
     try:
         info = get_collection_cluster_info(peer_url, collection_name)
     except:
         return False
 
     remote_shards = info["remote_shards"]
+    local_shards = info["local_shards"]
 
     if len(remote_shards) == 0:
         return False
 
-    return all(map(lambda shard: shard["state"] == "Active", remote_shards))
+    return all(map(lambda shard: shard["state"] == "Active", local_shards + remote_shards))