Files
qdrant/tests/openapi/test_shard_key.py
2025-11-28 21:50:09 -03:00

55 lines
1.7 KiB
Python

import os
import pytest
from .helpers.collection_setup import basic_collection_setup, drop_collection
from .helpers.helpers import request_with_validation
@pytest.fixture(autouse=True)
def setup(on_disk_vectors, collection_name):
basic_collection_setup(collection_name=collection_name, on_disk_vectors=on_disk_vectors, wal_capacity=1,
sharding_method="custom")
yield
drop_collection(collection_name=collection_name)
def create_shard_key(key, collection_name):
response = request_with_validation(
api='/collections/{collection_name}/shards',
method="PUT",
path_params={'collection_name': collection_name},
body={
"shard_key": key,
},
)
assert response.ok
def get_shard_keys(collection_name):
response = request_with_validation(
api='/collections/{collection_name}/shards',
method="GET",
path_params={'collection_name': collection_name},
)
assert response.ok
return response
@pytest.mark.skipif(
not os.getenv("QDRANT__CLUSTER__ENABLED"),
reason="only works in distributed mode"
)
def test_shard_keys_list(collection_name):
# no shard keys on collection
response = get_shard_keys(collection_name)
assert len(response.json()["result"]["shard_keys"]) == 0
create_shard_key("test_key", collection_name)
create_shard_key(100, collection_name)
response = get_shard_keys(collection_name)
shard_keys = response.json()["result"].get("shard_keys", [])
assert len(shard_keys) == 2
assert any("test_key" in item.values() for item in shard_keys)
assert any(100 in item.values() for item in shard_keys)