refactor configure index test

pinecone-io · Dec 6, 2023 · 9c6ccb6 · 9c6ccb6
1 parent f952fad
commit 9c6ccb6
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 65 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 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 = checkIfIndexExists(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 String createIndexIfNotExistsControlPlane(PineconeClientConfig config, int dimension) throws IOException, InterruptedException {
+        PineconeIndexOperationClient controlPlaneClient = new PineconeIndexOperationClient(config);
+        List<String> indexList = controlPlaneClient.listIndexes();
+        String indexName = checkIfIndexExists(indexList, dimension, controlPlaneClient);
+
+        return (indexName.isEmpty()) ? createNewIndex(dimension, controlPlaneClient) : indexName;
+    }
+
+    private static String checkIfIndexExists(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,13 +1,12 @@
-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.helpers.IndexManager;
 import io.pinecone.model.ConfigureIndexRequest;
-import io.pinecone.model.CreateIndexRequest;
 import io.pinecone.model.IndexMeta;
 import org.junit.jupiter.api.*;
 import org.slf4j.Logger;
@@ -19,30 +18,22 @@
 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 = new IndexManager().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
@@ -25,7 +25,7 @@ public class UpsertAndDeleteTest {
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = new IndexManager().createIndexIfNotExists(dimension);
+        PineconeConnection connection = new IndexManager().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
@@ -18,7 +18,7 @@ public class UpsertAndDescribeIndexStatsTest {
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = new IndexManager().createIndexIfNotExists(dimension);
+        PineconeConnection connection = new IndexManager().createIndexIfNotExistsDataPlane(dimension);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }