From 9c6ccb69ae709e305db184a6f88f0c4d4662db6b Mon Sep 17 00:00:00 2001 From: rohanshah18 Date: Wed, 6 Dec 2023 17:32:29 -0500 Subject: [PATCH] refactor configure index test --- ...eIndexOperationsClientIntegrationTest.java | 3 +- .../io/pinecone/helpers/IndexManager.java | 58 +++++++++++------ .../ConfigureIndexTest.java | 63 +++++++------------ .../dataplane/UpsertAndDeleteTest.java | 2 +- .../UpsertAndDescribeIndexStatsTest.java | 2 +- 5 files changed, 63 insertions(+), 65 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..e6c27731 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 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 = 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 indexList = controlPlaneClient.listIndexes(); + String indexName = checkIfIndexExists(indexList, dimension, controlPlaneClient); + + return (indexName.isEmpty()) ? createNewIndex(dimension, controlPlaneClient) : indexName; + } + + private static String checkIfIndexExists(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..ed40f38e 100644 --- a/src/integration/java/io/pinecone/integration/ConfigureIndexTest.java +++ b/src/integration/java/io/pinecone/integration/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 index 0642bfd5..036d5dfb 100644 --- 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 index c5165e09..8b229eab 100644 --- 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(); }