Refactor configure index test (#50)

## Problem After recent changes in integration test to use the existing index if it exists, the configure index test was required to be updated to consume less resources. ## Solution 1. Refactored checkIfIndexExists code: a. To support both data and control plane integration tests b. Check if any index matches the criteria compared to old approach of checking the first index of the list and otherwise create a new index. 2. For configure test: a. Updated the beforeAll logic to check if an index exist b. Deleted the afterEach block to delete each index and save them for future test ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [X] None of the above: (explain here) Added Integration tests ## Test Plan Ran integration tests locally.
pinecone-io · Dec 7, 2023 · 5426722 · 5426722
1 parent f952fad
commit 5426722
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 67 deletions.
diff --git a/src/integration/java/io/pinecone/PineconeIndexOperationsClientIntegrationTest.java b/src/integration/java/io/pinecone/PineconeIndexOperationsClientIntegrationTest.java
@@ -25,7 +25,7 @@ public void setUp() {
     }
 
     @Test
-    public void createAndDelete() throws IOException {
+    public void createAndDelete() throws IOException, InterruptedException {
         String indexName = RandomStringBuilder.build("index-name", 8);
 
         // Create an index
@@ -48,5 +48,6 @@ public void createAndDelete() throws IOException {
 
         // Cleanup
         pinecone.deleteIndex(indexName);
+        Thread.sleep(3500);
     }
 }
diff --git a/src/integration/java/io/pinecone/helpers/IndexManager.java b/src/integration/java/io/pinecone/helpers/IndexManager.java
@@ -8,31 +8,15 @@
 import java.util.List;
 
 public class IndexManager {
-    public PineconeConnection createIndexIfNotExists(int dimension) throws IOException, InterruptedException {
-        boolean createNewIndex = false;
-        String indexName = "";
+    public static PineconeConnection createIndexIfNotExistsDataPlane(int dimension) throws IOException, InterruptedException {
         PineconeClientConfig config = new PineconeClientConfig()
                 .withApiKey(System.getenv("PINECONE_API_KEY"))
                 .withEnvironment(System.getenv("PINECONE_ENVIRONMENT"));
         PineconeIndexOperationClient controlPlaneClient = new PineconeIndexOperationClient(config);
         List<String> indexList = controlPlaneClient.listIndexes();
 
-        if (!indexList.isEmpty()) {
-            indexName = indexList.get(0);
-            IndexMeta indexMeta = isIndexReady(indexName, controlPlaneClient);
-            if (indexMeta.getDatabase().getDimension() != dimension) {
-                createNewIndex = true;
-            }
-        }
-
-        if (createNewIndex) {
-            indexName = RandomStringBuilder.build("index-name", 8);
-            CreateIndexRequest createIndexRequest = new CreateIndexRequest()
-                    .withIndexName(indexName)
-                    .withDimension(dimension)
-                    .withMetric("euclidean");
-            controlPlaneClient.createIndex(createIndexRequest);
-        }
+        String indexName = findIndexWithDimensionAndPodType(indexList, dimension, controlPlaneClient);
+        if(indexName.isEmpty()) indexName = createNewIndex(dimension, controlPlaneClient);
 
         PineconeClient dataPlaneClient = new PineconeClient(config);
         IndexMeta indexMeta = controlPlaneClient.describeIndex(indexName);
@@ -43,6 +27,40 @@ public PineconeConnection createIndexIfNotExists(int dimension) throws IOExcepti
                         .withConnectionUrl("https://" + host));
     }
 
+    public static String createIndexIfNotExistsControlPlane(PineconeClientConfig config, int dimension) throws IOException, InterruptedException {
+        PineconeIndexOperationClient controlPlaneClient = new PineconeIndexOperationClient(config);
+        List<String> indexList = controlPlaneClient.listIndexes();
+        String indexName = findIndexWithDimensionAndPodType(indexList, dimension, controlPlaneClient);
+
+        return (indexName.isEmpty()) ? createNewIndex(dimension, controlPlaneClient) : indexName;
+    }
+
+    private static String findIndexWithDimensionAndPodType(List<String> indexList, int dimension, PineconeIndexOperationClient controlPlaneClient)
+            throws IOException, InterruptedException {
+        int i = 0;
+        while (i < indexList.size()) {
+            IndexMeta indexMeta = isIndexReady(indexList.get(i), controlPlaneClient);
+            if (indexMeta.getDatabase().getDimension() == dimension
+                    && indexMeta.getDatabase().getPodType().equals("p1.x1")) {
+                return indexList.get(i);
+            }
+            i++;
+        }
+        return "";
+    }
+
+    private static String createNewIndex(int dimension, PineconeIndexOperationClient controlPlaneClient) throws IOException {
+        String indexName = RandomStringBuilder.build("index-name", 8);
+        CreateIndexRequest createIndexRequest = new CreateIndexRequest()
+                .withIndexName(indexName)
+                .withDimension(dimension)
+                .withMetric("euclidean")
+                .withPodType("p1.x1");
+        controlPlaneClient.createIndex(createIndexRequest);
+
+        return indexName;
+    }
+
     public static IndexMeta isIndexReady(String indexName, PineconeIndexOperationClient indexOperationClient)
             throws IOException, InterruptedException {
         IndexMeta indexMeta;
@@ -51,9 +69,9 @@ public static IndexMeta isIndexReady(String indexName, PineconeIndexOperationCli
             if (indexMeta.getStatus().getState().equalsIgnoreCase("ready")) {
                 break;
             }
-
             Thread.sleep(1000);
         }
+
         return indexMeta;
     }
 }
diff --git a/...econe/integration/ConfigureIndexTest.java → ...tion/controlPlane/ConfigureIndexTest.java b/...econe/integration/ConfigureIndexTest.java → ...tion/controlPlane/ConfigureIndexTest.java
@@ -1,48 +1,39 @@
-package io.pinecone.integration;
+package io.pinecone.integration.controlPlane;
 
 import io.pinecone.PineconeClientConfig;
 import io.pinecone.PineconeClientLiveIntegTest;
 import io.pinecone.PineconeIndexOperationClient;
 import io.pinecone.exceptions.PineconeBadRequestException;
 import io.pinecone.exceptions.PineconeNotFoundException;
-import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.model.ConfigureIndexRequest;
-import io.pinecone.model.CreateIndexRequest;
 import io.pinecone.model.IndexMeta;
 import org.junit.jupiter.api.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 
+import static io.pinecone.helpers.IndexManager.createIndexIfNotExistsControlPlane;
 import static io.pinecone.helpers.IndexManager.isIndexReady;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class ConfigureIndexTest {
+    private static PineconeClientConfig config;
     private PineconeIndexOperationClient indexOperationClient;
     private String indexName;
     private static final Logger logger = LoggerFactory.getLogger(PineconeClientLiveIntegTest.class);
 
-    @BeforeEach
-    public void setUp() throws IOException {
-        indexName = RandomStringBuilder.build("index-name", 8);
-        PineconeClientConfig config = new PineconeClientConfig()
+    @BeforeAll
+    public static void defineConfig() {
+        config = new PineconeClientConfig()
                 .withApiKey(System.getenv("PINECONE_API_KEY"))
                 .withEnvironment(System.getenv("PINECONE_ENVIRONMENT"));
-        indexOperationClient = new PineconeIndexOperationClient(config);
-
-        // Create an index
-        CreateIndexRequest request = new CreateIndexRequest()
-                .withIndexName(indexName)
-                .withDimension(5)
-                .withMetric("euclidean");
-        indexOperationClient.createIndex(request);
     }
 
-    @AfterEach
-    public void cleanUp() throws IOException, InterruptedException {
-        indexOperationClient.deleteIndex(indexName);
-        Thread.sleep(3500);
+    @BeforeEach
+    public void setUp() throws IOException, InterruptedException {
+        indexName = createIndexIfNotExistsControlPlane(config, 5);
+        indexOperationClient = new PineconeIndexOperationClient(config);
     }
 
     @Test
@@ -76,28 +67,7 @@ public void configureIndexExceedingQuota() {
     }
 
     @Test
-    public void scaleUp() {
-        try{
-            // Verify the starting state
-            IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
-            assertEquals(1, indexMeta.getDatabase().getReplicas());
-
-            // Configure the index
-            ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
-                    .withReplicas(2);
-            isIndexReady(indexName, indexOperationClient);
-            indexOperationClient.configureIndex(indexName, configureIndexRequest);
-
-            // Verify replicas were scaled up
-            indexMeta = indexOperationClient.describeIndex(indexName);
-            assertEquals(2, indexMeta.getDatabase().getReplicas());
-        } catch (Exception exception) {
-            logger.error(exception.toString());
-        }
-    }
-
-    @Test
-    public void scaleDown() {
+    public void scaleUpAndDown() {
         try {
             // Verify the starting state
             IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
@@ -163,13 +133,17 @@ public void sizeIncrease() {
             // Get the index description to verify the new pod type
             indexMeta = indexOperationClient.describeIndex(indexName);
             assertEquals("p1.x2", indexMeta.getDatabase().getPodType());
+
+            // Delete this index since it'll be unused for future tests
+            indexOperationClient.deleteIndex(indexName);
+            Thread.sleep(3500);
         } catch (Exception exception) {
             logger.error(exception.getLocalizedMessage());
         }
     }
 
     @Test
-    public void sizeDown() {
+    public void sizeDown() throws IOException, InterruptedException {
         try {
             // Verify the starting state
             IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
@@ -189,9 +163,14 @@ public void sizeDown() {
             configureIndexRequest = new ConfigureIndexRequest()
                     .withPodType("p1.x1");
             indexOperationClient.configureIndex(indexName, configureIndexRequest);
+            Thread.sleep(3500);
         } catch (Exception exception) {
             assertEquals(exception.getClass(), PineconeBadRequestException.class);
             assertEquals(exception.getMessage(), "scaling down pod type is not supported");
+        } finally {
+            // Delete this index since it'll be unused for other tests
+            indexOperationClient.deleteIndex(indexName);
+            Thread.sleep(3500);
         }
     }
 }
diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java
@@ -3,7 +3,6 @@
 import com.google.protobuf.Struct;
 import com.google.protobuf.Value;
 import io.pinecone.PineconeConnection;
-import io.pinecone.helpers.IndexManager;
 import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.proto.*;
 import org.junit.jupiter.api.BeforeAll;
@@ -15,6 +14,7 @@
 import java.util.concurrent.ExecutionException;
 
 import static io.pinecone.helpers.BuildUpsertRequest.*;
+import static io.pinecone.helpers.IndexManager.createIndexIfNotExistsDataPlane;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -25,7 +25,7 @@ public class UpsertAndDeleteTest {
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = new IndexManager().createIndexIfNotExists(dimension);
+        PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }

diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java
@@ -1,11 +1,11 @@
 package io.pinecone.integration.dataplane;
 
 import io.pinecone.*;
-import io.pinecone.helpers.IndexManager;
 import io.pinecone.proto.*;
 import org.junit.jupiter.api.*;
 
 import static io.pinecone.helpers.BuildUpsertRequest.*;
+import static io.pinecone.helpers.IndexManager.createIndexIfNotExistsDataPlane;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.io.IOException;
@@ -18,7 +18,7 @@ public class UpsertAndDescribeIndexStatsTest {
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = new IndexManager().createIndexIfNotExists(dimension);
+        PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }