Fix unit tests

pinecone-io · Jul 18, 2024 · 1378390 · 1378390
1 parent 343751c
commit 1378390
Showing 1 changed file with 42 additions and 17 deletions.
diff --git a/tests/unit/models/test_index_list.py b/tests/unit/models/test_index_list.py
@@ -1,26 +1,41 @@
 import pytest
 from pinecone import IndexList
-from pinecone.core.openapi.control.models import IndexList as OpenApiIndexList
+from pinecone.core.openapi.control.models import (
+    IndexList as OpenApiIndexList,
+    IndexModel as OpenApiIndexModel,
+    IndexModelSpec,
+    IndexModelStatus,
+    DeletionProtection,
+    PodSpec as OpenApiPodSpec,
+)
 
 
 @pytest.fixture
 def index_list_response():
     return OpenApiIndexList(
         indexes=[
-            {
-                "name": "test-index-1",
-                "dimension": 2,
-                "metric": "cosine",
-                "spec": {
-                    "pod": {"environment": "us-west1-gcp", "pod_type": "p1.x1", "pods": 1, "replicas": 1, "shards": 1}
-                },
-            },
-            {
-                "name": "test-index-2",
-                "dimension": 3,
-                "metric": "cosine",
-                "spec": {"serverless": {"cloud": "aws", "region": "us-west-2"}},
-            },
+            OpenApiIndexModel(
+                name="test-index-1",
+                dimension=2,
+                metric="cosine",
+                host="https://test-index-1.pinecone.io",
+                status=IndexModelStatus(ready=True, state="Ready"),
+                deletion_protection=DeletionProtection("enabled"),
+                spec=IndexModelSpec(
+                    pod=OpenApiPodSpec(environment="us-west1-gcp", pod_type="p1.x1", pods=1, replicas=1, shards=1)
+                ),
+            ),
+            OpenApiIndexModel(
+                name="test-index-2",
+                dimension=3,
+                metric="cosine",
+                host="https://test-index-2.pinecone.io",
+                status=IndexModelStatus(ready=True, state="Ready"),
+                deletion_protection=DeletionProtection("disabled"),
+                spec=IndexModelSpec(
+                    pod=OpenApiPodSpec(environment="us-west1-gcp", pod_type="p1.x1", pods=1, replicas=1, shards=1)
+                ),
+            ),
         ],
         _check_type=False,
     )
@@ -42,8 +57,18 @@ def test_index_list_names_syntactic_sugar(self, index_list_response):
 
     def test_index_list_getitem(self, index_list_response):
         iil = IndexList(index_list_response)
-        assert iil[0] == index_list_response.indexes[0]
-        assert iil[1] == index_list_response.indexes[1]
+        input = index_list_response
+        assert input.indexes[0].name == iil[0].name
+        assert input.indexes[0].dimension == iil[0].dimension
+        assert input.indexes[0].metric == iil[0].metric
+        assert input.indexes[0].host == iil[0].host
+
+        # Deletion protection is a special case, since we're unwrapping
+        # the value from the DeletionProtection object inside our IndexModel
+        assert input.indexes[0].deletion_protection.value == iil[0].deletion_protection
+        assert iil[0].deletion_protection == "enabled"
+
+        assert input.indexes[1].name == iil[1].name
 
     def test_index_list_proxies_methods(self, index_list_response):
         # Forward compatibility, in case we add more attributes to IndexList for pagination