mirror of
https://github.com/qdrant/qdrant.git
synced 2026-08-06 10:00:58 -05:00
fc77e3e511
* Use systematically assert_http_ok for better error reporting * extraction assertion to use it outside of pytest
99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
import pathlib
|
|
|
|
from .utils import *
|
|
from .assertions import assert_http_ok
|
|
|
|
N_PEERS = 5
|
|
N_SHARDS = 4
|
|
N_REPLICA = 2
|
|
|
|
|
|
def test_collection_sharding(tmp_path: pathlib.Path):
|
|
assert_project_root()
|
|
peer_dirs = make_peer_folders(tmp_path, N_PEERS)
|
|
|
|
# Gathers REST API uris
|
|
peer_api_uris = []
|
|
|
|
# Start bootstrap
|
|
(bootstrap_api_uri, bootstrap_uri) = start_first_peer(
|
|
peer_dirs[0], "peer_0_0.log")
|
|
peer_api_uris.append(bootstrap_api_uri)
|
|
|
|
# Wait for leader
|
|
leader = wait_peer_added(bootstrap_api_uri)
|
|
|
|
# Start other peers
|
|
for i in range(1, len(peer_dirs)):
|
|
peer_api_uris.append(start_peer(
|
|
peer_dirs[i], f"peer_0_{i}.log", bootstrap_uri))
|
|
|
|
# Wait for cluster
|
|
wait_for_uniform_cluster_status(peer_api_uris, leader)
|
|
|
|
# Check that there are no collections on all peers
|
|
for uri in peer_api_uris:
|
|
r = requests.get(f"{uri}/collections")
|
|
assert_http_ok(r)
|
|
assert len(r.json()["result"]["collections"]) == 0
|
|
|
|
# Create collection in first peer
|
|
r = requests.put(
|
|
f"{peer_api_uris[0]}/collections/test_collection", json={
|
|
"vectors": {
|
|
"size": 4,
|
|
"distance": "Dot"
|
|
},
|
|
"shard_number": N_SHARDS,
|
|
"replication_factor": N_REPLICA,
|
|
})
|
|
assert_http_ok(r)
|
|
|
|
# Check that it exists on all peers
|
|
wait_collection_exists_and_active_on_all_peers(collection_name="test_collection", peer_api_uris=peer_api_uris)
|
|
|
|
# Check collection's cluster info
|
|
collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], "test_collection")
|
|
assert collection_cluster_info["shard_count"] == N_SHARDS
|
|
|
|
# Create points in first peer's collection
|
|
r = requests.put(
|
|
f"{peer_api_uris[0]}/collections/test_collection/points?wait=true", json={
|
|
"points": [
|
|
{
|
|
"id": 1,
|
|
"vector": [0.05, 0.61, 0.76, 0.74],
|
|
"payload": {
|
|
"city": "Berlin",
|
|
"country": "Germany",
|
|
"count": 1000000,
|
|
"square": 12.5,
|
|
"coords": {"lat": 1.0, "lon": 2.0}
|
|
}
|
|
},
|
|
{"id": 2, "vector": [0.19, 0.81, 0.75, 0.11],
|
|
"payload": {"city": ["Berlin", "London"]}},
|
|
{"id": 3, "vector": [0.36, 0.55, 0.47, 0.94],
|
|
"payload": {"city": ["Berlin", "Moscow"]}},
|
|
{"id": 4, "vector": [0.18, 0.01, 0.85, 0.80],
|
|
"payload": {"city": ["London", "Moscow"]}},
|
|
{"id": 5, "vector": [0.24, 0.18, 0.22, 0.44],
|
|
"payload": {"count": [0]}},
|
|
{"id": 6, "vector": [0.35, 0.08, 0.11, 0.44]}
|
|
]
|
|
})
|
|
assert_http_ok(r)
|
|
|
|
# Check that 'search' returns the same results on all peers
|
|
for uri in peer_api_uris:
|
|
r = requests.post(
|
|
f"{uri}/collections/test_collection/points/search", json={
|
|
"vector": [0.2, 0.1, 0.9, 0.7],
|
|
"top": 3,
|
|
}
|
|
)
|
|
assert_http_ok(r)
|
|
assert r.json()["result"][0]["id"] == 4
|
|
assert r.json()["result"][1]["id"] == 1
|
|
assert r.json()["result"][2]["id"] == 3
|