From 3ac374be74c50294b8bad21e6f3ca383043cc4e7 Mon Sep 17 00:00:00 2001 From: Tibor Reiss Date: Thu, 29 Aug 2024 20:29:11 +0200 Subject: [PATCH] Fix tests, annotations, signatures based on code review --- integration/test_cluster.py | 159 ++++++++++++++------------ weaviate/collections/cluster/sync.pyi | 6 + 2 files changed, 89 insertions(+), 76 deletions(-) diff --git a/integration/test_cluster.py b/integration/test_cluster.py index e29b50b8f..87a313502 100644 --- a/integration/test_cluster.py +++ b/integration/test_cluster.py @@ -1,4 +1,5 @@ -import pytest +from contextlib import contextmanager +from typing import Generator, List import weaviate from weaviate.collections.classes.config import ( @@ -6,90 +7,96 @@ DataType, Property, ) -from weaviate.util import parse_version_string +COLLECTION_NAME_PREFIX = "Collection_test_cluster" NODE_NAME = "node1" NUM_OBJECT = 10 -@pytest.fixture(scope="module") -def client(): +@contextmanager +def get_weaviate_client( + collection_names: List[str], +) -> Generator[weaviate.WeaviateClient, None, None]: client = weaviate.connect_to_local() - client.collections.delete_all() + for collection_name in collection_names: + client.collections.delete(collection_name) yield client - client.collections.delete_all() + for collection_name in collection_names: + client.collections.delete(collection_name) + client.close() -def test_rest_nodes_without_data(client: weaviate.WeaviateClient): +def test_rest_nodes_without_data() -> None: """get nodes status without data""" - resp = client.cluster.rest_nodes(output="verbose") - assert len(resp) == 1 - assert "gitHash" in resp[0] - assert resp[0]["name"] == NODE_NAME - assert resp[0]["shards"] is None - assert resp[0]["stats"]["objectCount"] == 0 - assert resp[0]["stats"]["shardCount"] == 0 - assert resp[0]["status"] == "HEALTHY" - assert "version" in resp[0] - - -def test_rest_nodes_with_data(client: weaviate.WeaviateClient): + with get_weaviate_client([]) as client: + resp = client.cluster.rest_nodes(output="verbose") + assert len(resp) == 1 + assert "gitHash" in resp[0] + assert resp[0]["name"] == NODE_NAME + assert resp[0]["shards"] is None + assert resp[0]["stats"]["objectCount"] == 0 + assert resp[0]["stats"]["shardCount"] == 0 + assert resp[0]["status"] == "HEALTHY" + assert "version" in resp[0] + + +def test_rest_nodes_with_data() -> None: """get nodes status with data""" - collection_name_1 = "Collection_1" - uncap_collection_name_1 = "collection_1" - collection = client.collections.create( - name=collection_name_1, - properties=[Property(name="Name", data_type=DataType.TEXT)], - vectorizer_config=Configure.Vectorizer.none(), - ) - collection.data.insert_many([{"Name": f"name {i}"} for i in range(NUM_OBJECT)]) - - collection_name_2 = "Collection_2" - collection = client.collections.create( - name=collection_name_2, - properties=[Property(name="Name", data_type=DataType.TEXT)], - vectorizer_config=Configure.Vectorizer.none(), - ) - collection.data.insert_many([{"Name": f"name {i}"} for i in range(NUM_OBJECT * 2)]) - - # server behaviour changed by https://github.com/weaviate/weaviate/pull/4203 - server_is_at_least_124 = parse_version_string( - client.get_meta()["version"] - ) > parse_version_string("1.24") - - resp = client.cluster.rest_nodes(output="verbose") - assert len(resp) == 1 - assert "gitHash" in resp[0] - assert resp[0]["name"] == NODE_NAME - assert len(resp[0]["shards"]) == 2 - assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT * 3 - assert resp[0]["stats"]["shardCount"] == 2 - assert resp[0]["status"] == "HEALTHY" - assert "version" in resp[0] - - shards = sorted(resp[0]["shards"], key=lambda x: x["class"]) - assert shards[0]["class"] == collection_name_1 - assert shards[0]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT - assert shards[1]["class"] == collection_name_2 - assert shards[1]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT * 2 - - resp = client.cluster.rest_nodes(collection=collection_name_1, output="verbose") - assert len(resp) == 1 - assert "gitHash" in resp[0] - assert resp[0]["name"] == NODE_NAME - assert len(resp[0]["shards"]) == 1 - assert resp[0]["stats"]["shardCount"] == 1 - assert resp[0]["status"] == "HEALTHY" - assert "version" in resp[0] - assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT - - resp = client.cluster.rest_nodes(uncap_collection_name_1, output="verbose") - assert len(resp) == 1 - assert "gitHash" in resp[0] - assert resp[0]["name"] == NODE_NAME - assert len(resp[0]["shards"]) == 1 - assert resp[0]["stats"]["shardCount"] == 1 - assert resp[0]["status"] == "HEALTHY" - assert "version" in resp[0] - assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT + collection_name_1 = f"{COLLECTION_NAME_PREFIX}_rest_nodes_with_data_1" + collection_name_2 = f"{COLLECTION_NAME_PREFIX}_rest_nodes_with_data_2" + uncap_collection_name_1 = collection_name_1[0].lower() + collection_name_1[1:] + + with get_weaviate_client([collection_name_1, collection_name_2]) as client: + collection = client.collections.create( + name=collection_name_1, + properties=[Property(name="Name", data_type=DataType.TEXT)], + vectorizer_config=Configure.Vectorizer.none(), + ) + collection.data.insert_many([{"Name": f"name {i}"} for i in range(NUM_OBJECT)]) + + collection = client.collections.create( + name=collection_name_2, + properties=[Property(name="Name", data_type=DataType.TEXT)], + vectorizer_config=Configure.Vectorizer.none(), + ) + collection.data.insert_many([{"Name": f"name {i}"} for i in range(NUM_OBJECT * 2)]) + + # server behaviour changed by https://github.com/weaviate/weaviate/pull/4203 + server_is_at_least_124 = client._connection._weaviate_version.is_at_least(1, 24, 0) + + resp = client.cluster.rest_nodes(output="verbose") + assert len(resp) == 1 + assert "gitHash" in resp[0] + assert resp[0]["name"] == NODE_NAME + assert resp[0]["shards"] is not None and len(resp[0]["shards"]) == 2 + assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT * 3 + assert resp[0]["stats"]["shardCount"] == 2 + assert resp[0]["status"] == "HEALTHY" + assert "version" in resp[0] + + shards = sorted(resp[0]["shards"], key=lambda x: x["class"]) + assert shards[0]["class"] == collection_name_1 + assert shards[0]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT + assert shards[1]["class"] == collection_name_2 + assert shards[1]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT * 2 + + resp = client.cluster.rest_nodes(collection=collection_name_1, output="verbose") + assert len(resp) == 1 + assert "gitHash" in resp[0] + assert resp[0]["name"] == NODE_NAME + assert resp[0]["shards"] is not None and len(resp[0]["shards"]) == 1 + assert resp[0]["stats"]["shardCount"] == 1 + assert resp[0]["status"] == "HEALTHY" + assert "version" in resp[0] + assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT + + resp = client.cluster.rest_nodes(uncap_collection_name_1, output="verbose") + assert len(resp) == 1 + assert "gitHash" in resp[0] + assert resp[0]["name"] == NODE_NAME + assert resp[0]["shards"] is not None and len(resp[0]["shards"]) == 1 + assert resp[0]["stats"]["shardCount"] == 1 + assert resp[0]["status"] == "HEALTHY" + assert "version" in resp[0] + assert resp[0]["stats"]["objectCount"] == 0 if server_is_at_least_124 else NUM_OBJECT diff --git a/weaviate/collections/cluster/sync.pyi b/weaviate/collections/cluster/sync.pyi index a32e8ff62..76da2e310 100644 --- a/weaviate/collections/cluster/sync.pyi +++ b/weaviate/collections/cluster/sync.pyi @@ -1,5 +1,6 @@ from typing import List, Literal, Optional, overload +from weaviate.cluster.types import Node as NodeREST from weaviate.collections.classes.cluster import Node, Shards, Stats from weaviate.collections.cluster.cluster import _ClusterBase @@ -25,3 +26,8 @@ class _Cluster(_ClusterBase): *, output: Literal["verbose"], ) -> List[Node[Shards, Stats]]: ... + def rest_nodes( + self, + collection: Optional[str] = None, + output: Optional[Literal["minimal", "verbose"]] = None, + ) -> List[NodeREST]: ...