Cleanup test resources

Author: jhamon

scripts/repl.py

@@ -2,6 +2,7 @@
 import code
 from pinecone import Pinecone
 import logging
+import os
 
 
 def main():
@@ -15,18 +16,66 @@ def main():
     logging.basicConfig(
         level=logging.DEBUG, format="%(levelname)-8s | %(name)s:%(lineno)d | %(message)s"
     )
+    logger = logging.getLogger(__name__)
 
     # Start the interactive REPL
     banner = """
     Welcome to the custom Python REPL!
     Your initialization steps have been completed.
+
+    Two Pinecone objects are available:
+    - pc: Interact with the one-offs project
+    - pcci: Interact with the pinecone-python-client project (CI testing)
+
+    You can use the following functions to clean up the environment:
+    - delete_all_indexes(pc)
+    - delete_all_collections(pc)
+    - delete_all_backups(pc)
+    - cleanup_all(pc)
     """
 
+    def delete_all_indexes(pc):
+        for index in pc.db.index.list():
+            logger.info(f"Deleting index {index.name}")
+            try:
+                if index.deletion_protection == "enabled":
+                    logger.info(f"Disabling deletion protection for index {index.name}")
+                    pc.db.index.configure(name=index.name, deletion_protection="disabled")
+                pc.db.index.delete(name=index.name)
+            except Exception as e:
+                logger.error(f"Error deleting index {index.name}: {e}")
+
+    def delete_all_collections(pc):
+        for collection in pc.db.collection.list():
+            logger.info(f"Deleting collection {collection.name}")
+            try:
+                pc.db.collection.delete(name=collection.name)
+            except Exception as e:
+                logger.error(f"Error deleting collection {collection.name}: {e}")
+
+    def delete_all_backups(pc):
+        for backup in pc.db.backup.list():
+            logger.info(f"Deleting backup {backup.name}")
+            try:
+                pc.db.backup.delete(backup_id=backup.backup_id)
+            except Exception as e:
+                logger.error(f"Error deleting backup {backup.name}: {e}")
+
+    def cleanup_all(pc):
+        delete_all_indexes(pc)
+        delete_all_collections(pc)
+        delete_all_backups(pc)
+
     # Create a custom namespace with any pre-loaded variables
     namespace = {
         "__name__": "__main__",
         "__doc__": None,
         "pc": Pinecone(),
+        "pcci": Pinecone(api_key=os.environ.get("PINECONE_API_KEY_CI_TESTING")),
+        "delete_all_indexes": delete_all_indexes,
+        "delete_all_collections": delete_all_collections,
+        "delete_all_backups": delete_all_backups,
+        "cleanup_all": cleanup_all,
         # Add any other variables you want to have available in the REPL
     }
 

tests/integration/control/resources/index/test_create.py

@@ -17,11 +17,12 @@
 
 
 class TestCreateServerlessIndexHappyPath:
-    def test_create_index(self, pc: Pinecone, index_name):
+    def test_create_index(self, pc: Pinecone, index_name, index_tags):
         resp = pc.db.index.create(
             name=index_name,
             dimension=10,
             spec=ServerlessSpec(cloud=CloudProvider.AWS, region=AwsRegion.US_EAST_1),
+            tags=index_tags,
         )
         assert resp.name == index_name
         assert resp.dimension == 10
@@ -36,23 +37,25 @@ def test_create_index(self, pc: Pinecone, index_name):
         assert desc.deletion_protection == "disabled"  # default value
         assert desc.vector_type == "dense"  # default value
 
-    def test_create_skip_wait(self, pc, index_name):
+    def test_create_skip_wait(self, pc, index_name, index_tags):
         resp = pc.db.index.create(
             name=index_name,
             dimension=10,
             spec=ServerlessSpec(cloud=CloudProvider.AWS, region=AwsRegion.US_EAST_1),
             timeout=-1,
+            tags=index_tags,
         )
         assert resp.name == index_name
         assert resp.dimension == 10
         assert resp.metric == "cosine"
 
-    def test_create_infinite_wait(self, pc, index_name):
+    def test_create_infinite_wait(self, pc, index_name, index_tags):
         resp = pc.db.index.create(
             name=index_name,
             dimension=10,
             spec=ServerlessSpec(cloud=CloudProvider.AWS, region=AwsRegion.US_EAST_1),
             timeout=None,
+            tags=index_tags,
         )
         assert resp.name == index_name
         assert resp.dimension == 10
@@ -70,23 +73,23 @@ def test_create_default_index_with_metric(self, pc, create_index_params, metric)
         assert desc.vector_type == "dense"
 
     @pytest.mark.parametrize(
-        "metric_enum,vector_type_enum,dim,tags",
+        "metric_enum,vector_type_enum,dim",
         [
-            (Metric.COSINE, VectorType.DENSE, 10, None),
-            (Metric.EUCLIDEAN, VectorType.DENSE, 10, {"env": "prod"}),
-            (Metric.DOTPRODUCT, VectorType.SPARSE, None, {"env": "dev"}),
+            (Metric.COSINE, VectorType.DENSE, 10),
+            (Metric.EUCLIDEAN, VectorType.DENSE, 10),
+            (Metric.DOTPRODUCT, VectorType.SPARSE, None),
         ],
     )
     def test_create_with_enum_values(
-        self, pc, index_name, metric_enum, vector_type_enum, dim, tags
+        self, pc, index_name, metric_enum, vector_type_enum, dim, index_tags
     ):
         args = {
             "name": index_name,
             "metric": metric_enum,
             "vector_type": vector_type_enum,
             "deletion_protection": DeletionProtection.DISABLED,
             "spec": ServerlessSpec(cloud=CloudProvider.AWS, region=AwsRegion.US_EAST_1),
-            "tags": tags,
+            "tags": index_tags,
         }
         if dim is not None:
             args["dimension"] = dim
@@ -101,8 +104,7 @@ def test_create_with_enum_values(
         assert desc.name == index_name
         assert desc.spec.serverless.cloud == "aws"
         assert desc.spec.serverless.region == "us-east-1"
-        if tags:
-            assert desc.tags.to_dict() == tags
+        assert desc.tags.to_dict() == index_tags
 
     @pytest.mark.parametrize("metric", ["cosine", "euclidean", "dotproduct"])
     def test_create_dense_index_with_metric(self, pc, create_index_params, metric):
@@ -113,13 +115,6 @@ def test_create_dense_index_with_metric(self, pc, create_index_params, metric):
         assert desc.metric == metric
         assert desc.vector_type == "dense"
 
-    def test_create_with_optional_tags(self, pc, create_index_params):
-        tags = {"foo": "FOO", "bar": "BAR"}
-        create_index_params["tags"] = tags
-        pc.db.index.create(**create_index_params)
-        desc = pc.db.index.describe(create_index_params["name"])
-        assert desc.tags.to_dict() == tags
-
 
 class TestCreatePodIndexHappyPath:
     def test_create_index_minimal_config(
@@ -152,7 +147,7 @@ def test_create_index_with_spec_options(
             metric="cosine",
             spec=PodSpec(
                 environment=pod_environment,
-                pod_type="p1.x2",
+                pod_type="p1.x1",
                 replicas=2,
                 metadata_config={"indexed": ["foo", "bar"]},
             ),
@@ -164,9 +159,10 @@ def test_create_index_with_spec_options(
         assert desc.dimension == 10
         assert desc.metric == "cosine"
         assert desc.spec.pod.environment == pod_environment
-        assert desc.spec.pod.pod_type == "p1.x2"
+        assert desc.spec.pod.pod_type == "p1.x1"
         assert desc.spec.pod.replicas == 2
         assert desc.spec.pod.metadata_config.indexed == ["foo", "bar"]
+        assert desc.tags.to_dict() == index_tags
 
     def test_create_index_with_deletion_protection(
         self, pc: Pinecone, index_name, pod_environment, index_tags

tests/integration/control_asyncio/resources/index/test_create.py

@@ -13,11 +13,11 @@
 @pytest.mark.asyncio
 class TestAsyncioCreateIndex:
     @pytest.mark.parametrize("spec_fixture", ("spec1", "spec2", "spec3"))
-    async def test_create_index(self, index_name, request, spec_fixture):
+    async def test_create_index(self, index_name, request, spec_fixture, index_tags):
         pc = PineconeAsyncio()
         spec = request.getfixturevalue(spec_fixture)
 
-        resp = await pc.db.index.create(name=index_name, dimension=10, spec=spec)
+        resp = await pc.db.index.create(name=index_name, dimension=10, spec=spec, tags=index_tags)
 
         assert resp.name == index_name
         assert resp.dimension == 10
@@ -33,17 +33,21 @@ async def test_create_index(self, index_name, request, spec_fixture):
         assert desc.vector_type == "dense"  # default value
         await pc.close()
 
-    async def test_create_skip_wait(self, index_name, spec1):
+    async def test_create_skip_wait(self, index_name, spec1, index_tags):
         pc = PineconeAsyncio()
-        resp = await pc.db.index.create(name=index_name, dimension=10, spec=spec1, timeout=-1)
+        resp = await pc.db.index.create(
+            name=index_name, dimension=10, spec=spec1, timeout=-1, tags=index_tags
+        )
         assert resp.name == index_name
         assert resp.dimension == 10
         assert resp.metric == "cosine"
         await pc.close()
 
-    async def test_create_infinite_wait(self, index_name, spec1):
+    async def test_create_infinite_wait(self, index_name, spec1, index_tags):
         async with PineconeAsyncio() as pc:
-            resp = await pc.db.index.create(name=index_name, dimension=10, spec=spec1, timeout=None)
+            resp = await pc.db.index.create(
+                name=index_name, dimension=10, spec=spec1, timeout=None, tags=index_tags
+            )
             assert resp.name == index_name
             assert resp.dimension == 10
             assert resp.metric == "cosine"
@@ -104,7 +108,7 @@ async def test_create_with_enum_values_and_tags(
     async def test_create_dense_index_with_metric(self, index_name, spec1, metric, index_tags):
         pc = PineconeAsyncio()
 
-        await pc.create_index(
+        await pc.db.index.create(
             name=index_name,
             dimension=10,
             spec=spec1,
@@ -121,7 +125,7 @@ async def test_create_dense_index_with_metric(self, index_name, spec1, metric, i
     async def test_create_with_optional_tags(self, index_name, spec1, index_tags):
         pc = PineconeAsyncio()
 
-        await pc.create_index(name=index_name, dimension=10, spec=spec1, tags=index_tags)
+        await pc.db.index.create(name=index_name, dimension=10, spec=spec1, tags=index_tags)
 
         desc = await pc.db.index.describe(name=index_name)
         assert desc.tags.to_dict() == index_tags
@@ -131,7 +135,7 @@ async def test_create_with_optional_tags(self, index_name, spec1, index_tags):
     async def test_create_sparse_index(self, index_name, spec1, index_tags):
         pc = PineconeAsyncio()
 
-        await pc.create_index(
+        await pc.db.index.create(
             name=index_name,
             spec=spec1,
             metric=Metric.DOTPRODUCT,
@@ -149,7 +153,7 @@ async def test_create_sparse_index(self, index_name, spec1, index_tags):
     async def test_create_with_deletion_protection(self, index_name, spec1, index_tags):
         pc = PineconeAsyncio()
 
-        await pc.create_index(
+        await pc.db.index.create(
             name=index_name,
             spec=spec1,
             metric=Metric.DOTPRODUCT,

tests/unit/test_rest_urllib3.py

@@ -0,0 +1,104 @@
+import pytest
+from unittest.mock import patch, MagicMock
+import urllib3
+from urllib3.exceptions import MaxRetryError
+from pinecone.openapi_support.rest_urllib3 import Urllib3RestClient
+from pinecone.config.openapi_configuration import Configuration
+
+
+class TestUrllib3RestClient:
+    @pytest.fixture
+    def config(self):
+        return Configuration(api_key="test-key")
+
+    @pytest.fixture
+    def client(self, config):
+        return Urllib3RestClient(config)
+
+    def test_retry_on_500_error(self, client):
+        # Mock response that fails with 500
+        mock_response = MagicMock()
+        mock_response.status = 500
+        mock_response.data = b'{"error": "Internal Server Error"}'
+        mock_response.headers = {}
+        mock_response.reason = "Internal Server Error"
+
+        # Mock pool manager to fail twice then succeed
+        with patch.object(client.pool_manager, "request") as mock_request:
+            mock_request.side_effect = [
+                urllib3.exceptions.HTTPError(response=mock_response),
+                urllib3.exceptions.HTTPError(response=mock_response),
+                mock_response,  # Success on third try
+            ]
+
+            # Make request
+            response = client.request(
+                method="GET",
+                url="https://api.pinecone.io/test",
+                headers={"Authorization": "test-key"},
+            )
+
+            # Verify request was made 3 times (initial + 2 retries)
+            assert mock_request.call_count == 3
+
+            # Verify the response is successful
+            assert response.status == 200
+
+    def test_max_retries_exceeded(self, client):
+        # Mock response that always fails with 500
+        mock_response = MagicMock()
+        mock_response.status = 500
+        mock_response.data = b'{"error": "Internal Server Error"}'
+        mock_response.headers = {}
+        mock_response.reason = "Internal Server Error"
+
+        # Mock pool manager to always fail
+        with patch.object(client.pool_manager, "request") as mock_request:
+            mock_request.side_effect = urllib3.exceptions.HTTPError(response=mock_response)
+
+            # Make request and expect MaxRetryError
+            with pytest.raises(MaxRetryError):
+                client.request(
+                    method="GET",
+                    url="https://api.pinecone.io/test",
+                    headers={"Authorization": "test-key"},
+                )
+
+            # Verify request was made 4 times (initial + 3 retries)
+            assert mock_request.call_count == 4
+
+    def test_custom_retry_config(self):
+        # Create custom retry configuration
+        custom_retry = urllib3.Retry(
+            total=2, backoff_factor=0.5, status_forcelist=(500, 502, 503, 504)
+        )
+
+        config = Configuration(api_key="test-key", retries=custom_retry)
+        client = Urllib3RestClient(config)
+
+        # Mock response that fails with 500
+        mock_response = MagicMock()
+        mock_response.status = 500
+        mock_response.data = b'{"error": "Internal Server Error"}'
+        mock_response.headers = {}
+        mock_response.reason = "Internal Server Error"
+
+        # Mock pool manager to fail once then succeed
+        with patch.object(client.pool_manager, "request") as mock_request:
+            mock_request.side_effect = [
+                urllib3.exceptions.HTTPError(response=mock_response),
+                mock_response,  # Success on second try
+            ]
+
+            # Make request
+            response = client.request(
+                method="GET",
+                url="https://api.pinecone.io/test",
+                headers={"Authorization": "test-key"},
+            )
+
+            # Verify request was made 2 times (initial + 1 retry)
+            assert mock_request.call_count == 2
+
+            # Verify the response is successful
+            assert response.status == 200