qdrant/tests/consensus_tests/test_remove_node.py

import pathlib

from .utils import *
from .assertions import assert_http_ok

N_PEERS = 3
N_SHARDS = 4
N_REPLICA = 2


def upsert_points(peer_url, city):
    # Create points in first peer's collection
    r_batch = requests.put(
        f"{peer_url}/collections/test_collection/points?wait=true", json={
            "points": [
                {"id": 1, "vector": [0.05, 0.61, 0.76, 0.74], "payload": {"city": city}},
                {"id": 2, "vector": [0.19, 0.81, 0.75, 0.11], "payload": {"city": city}},
                {"id": 3, "vector": [0.36, 0.55, 0.47, 0.94], "payload": {"city": city}},
                {"id": 4, "vector": [0.18, 0.01, 0.85, 0.80], "payload": {"city": city}},
                {"id": 5, "vector": [0.24, 0.18, 0.22, 0.44], "payload": {"city": city}},
                {"id": 6, "vector": [0.35, 0.08, 0.11, 0.44], "payload": {"city": city}},
                {"id": 7, "vector": [0.45, 0.07, 0.21, 0.04], "payload": {"city": city}},
                {"id": 8, "vector": [0.75, 0.18, 0.91, 0.48], "payload": {"city": city}},
                {"id": 9, "vector": [0.30, 0.01, 0.10, 0.12], "payload": {"city": city}},
                {"id": 10, "vector": [0.95, 0.8, 0.17, 0.19], "payload": {"city": city}},
            ]
        })
    assert_http_ok(r_batch)


def create_collection(peer_url, collection="test_collection"):
    # Create collection in first peer
    r_batch = requests.put(
        f"{peer_url}/collections/{collection}", json={
            "vectors": {
                "size": 4,
                "distance": "Dot"
            },
            "shard_number": N_SHARDS,
            "replication_factor": N_REPLICA,
        })
    assert_http_ok(r_batch)


def search(peer_url, city):
    q = {
        "vector": [0.2, 0.1, 0.9, 0.7],
        "top": 10,
        "with_vector": False,
        "with_payload": True,
        "filter": {
            "must": [
                {
                    "key": "city",
                    "match": {"value": city}
                }
            ]
        }
    }

    result = requests.post(f"{peer_url}/collections/test_collection/points/search", json=q).json()
    assert "result" in result, result
    assert result["result"] is not None, result
    return result["result"]


def test_remove_node(tmp_path: pathlib.Path):
    assert_project_root()

    peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS)

    create_collection(peer_api_uris[0])
    wait_collection_exists_and_active_on_all_peers(collection_name="test_collection", peer_api_uris=peer_api_uris)

    upsert_points(peer_api_uris[0], "Paris")

    search_result = search(peer_api_uris[0], "Paris")
    assert len(search_result) > 0

    # Kill last peer

    p = processes.pop()
    p.kill()

    # Validate search works with the dead node
    search_result = search(peer_api_uris[0], "Paris")
    assert len(search_result) > 0

    # Validate upsert works with the dead node
    upsert_points(peer_api_uris[0], "Berlin")

    # Assert that there are dead replicas
    wait_for_some_replicas_not_active(peer_api_uris[0], "test_collection")

    # Assert all records were changed
    search_result = search(peer_api_uris[0], "Paris")
    assert len(search_result) == 0

    collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], "test_collection")
    dead_peer = [shard["peer_id"] for shard in collection_cluster_info["remote_shards"] if shard['state'] != 'Active']

    assert len(dead_peer) > 1

    # force remove dead peer
    res = requests.delete(f"{peer_api_uris[0]}/cluster/peer/{dead_peer[0]}?force=true")
    assert_http_ok(res)

    # Assert all records were changed
    search_result = search(peer_api_uris[0], "Paris")
    assert len(search_result) == 0

    # Assert all records were changed
    search_result = search(peer_api_uris[0], "Berlin")
    assert len(search_result) > 0