From 5426722b52375513ceade8c64ddc0c70e70e989f Mon Sep 17 00:00:00 2001 From: Rohan Shah Date: Thu, 7 Dec 2023 14:51:42 -0500 Subject: [PATCH] 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. --- ...eIndexOperationsClientIntegrationTest.java | 3 +- .../io/pinecone/helpers/IndexManager.java | 58 +++++++++++------ .../ConfigureIndexTest.java | 63 +++++++------------ .../dataplane/UpsertAndDeleteTest.java | 4 +- .../UpsertAndDescribeIndexStatsTest.java | 4 +- 5 files changed, 65 insertions(+), 67 deletions(-) rename src/integration/java/io/pinecone/integration/{ => controlPlane}/ConfigureIndexTest.java (80%) diff --git a/src/integration/java/io/pinecone/PineconeIndexOperationsClientIntegrationTest.java b/src/integration/java/io/pinecone/PineconeIndexOperationsClientIntegrationTest.java index 06e8078b..071be2a3 100644 --- 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 index 3071c964..046b9266 100644 --- 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 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 indexList = controlPlaneClient.listIndexes(); + String indexName = findIndexWithDimensionAndPodType(indexList, dimension, controlPlaneClient); + + return (indexName.isEmpty()) ? createNewIndex(dimension, controlPlaneClient) : indexName; + } + + private static String findIndexWithDimensionAndPodType(List 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/src/integration/java/io/pinecone/integration/ConfigureIndexTest.java b/src/integration/java/io/pinecone/integration/controlPlane/ConfigureIndexTest.java similarity index 80% rename from src/integration/java/io/pinecone/integration/ConfigureIndexTest.java rename to src/integration/java/io/pinecone/integration/controlPlane/ConfigureIndexTest.java index e498c63b..d945c5d2 100644 --- a/src/integration/java/io/pinecone/integration/ConfigureIndexTest.java +++ b/src/integration/java/io/pinecone/integration/controlPlane/ConfigureIndexTest.java @@ -1,13 +1,11 @@ -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; @@ -15,34 +13,27 @@ 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 index 0642bfd5..ad87054e 100644 --- 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 index c5165e09..98668a55 100644 --- 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(); }