From 7bb00eb27e8e477752720416614c4aea27cb8637 Mon Sep 17 00:00:00 2001
From: Tibor Reiss <tibor.reiss@gmail.com>
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]: ...