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(); }