Refactor CollectionTest, ConfigureIndexTest, and IndexManager to improve integration test speed and reliability (#81)

## Problem
We've had continued issues with flakiness of integration tests, and
tests taking a while to run (25 - 30 minutes for the entire `pr`
workflow in some cases). A lot of this is waiting for indexes and
collections to be `Ready` instead of `Initializing`. Specifically, the
`CollectionTest` and `CollectionErrorTest` can take a really long time
due to setting up and waiting on multiple indexes and collections.

While working on the above issue, I also noticed test failures and
repeated issues due to `findIndexWithDimensionAndType`, and tests
running with indexes created by and belonging to other test runs. The
`isIndexReady` function also tends to be inefficient as the polling
isn't consistent, and we can end up waiting sometimes minutes longer
than we may need to.

With the new `Pinecone` and `Index`/`AsyncIndex` class structure in
place, we can refactor our integration tests around the new pattern.

## Solution
Ultimately, I would like to move our tests away from the
`createIndexIfNotExistsDataPlane` / `createIndexIfNotExistsControlPlane`
functions and their reliance on the `findIndexWithDimensionAndType`
function, which should be deprecated. We should come up with a pattern
for sharing resources across all tests in a run, and setting up /
tearing those down once. This can be worked on in subsequent PRs.

- Move `CollectionErrorTests` tests into `CollectionTest`. Share index
and collection setup `@BeforeAll` across collections tests. Wait less
time for the indexes created from the collection to be ready as this
specifically can take a number of minutes. Relax some of the assertions
on the `status` of created collections and indexes as I don't think we
need to be that thorough here.
- Update `createIndexIfNotExistsDataPlane` to return a tuple
(`AbstractMap.SimpleEntry<String, Pinecone>`) of the `indexName` and the
`Pinecone` client to make things a bit easier to work with.
- Update all of the `dataPlane/` tests to use
`pineconeClient.createIndexConnection()` and
`pineconeClient.createAsyncIndexConnection()` for managing data plane
operations.
- Deprecate `isIndexReady` in lieu of `waitUntilIndexIsReady`. The
polling is more consistent with this function, and it most likely saves
us time overall.
- Refactor `ConfigureIndexTest` to create its own index and clean up
after.
- Fix issue with `findIndexByDimensionAndType` calling `isIndexReady()`
on each index while iterating through the list.
- Fix issue in `IndexInterface.validateUpsertRequest` where we were
trying to call
`sparseValuesWithUnsignedIndices.getIndicesWithUnsigned32IntList()` and
`sparseValuesWithUnsignedIndices.getValuesList()` on a possible `null`
causing a `NullPointerException`.
- Talked with @ssmith-pc, and I think it's standard practice in tests to
not `try`/`catch` yourself unless you need to assert on the result. We
should be letting errors throw to the test runner and let it handle them
so we're not clobbering logs and stack traces. I've cleaned up
`try`/`catch` statements which don't seem to be needed.
- Running `gradle integrationTest --info` in `pr.yml` to get more
detailed log output in the console for better troubleshooting of ongoing
flapping.
- Adding `assertWithRetry` wrappers for specific actions which have been
troublesome. Adding `Thread.sleep()` to a few places to avoid hammering
an index too quickly / etc.
- Added a new `describeIndexStats()` overload to `IndexInterface` and
`Index` / `AsyncIndex` to allow calling without needing to explicitly
pass `null`.
 
I spent a lot of time running these locally and in CI to see how they
perform. Overall, it seems like these changes improve overall
reliability, although we still do see a few failures like the gRPC `no
healthy upstream` on data plane operations with fresh indexes.

The total amount of time it takes to run both sets of integration tests
has been cut significantly in most cases:

<img width="1130" alt="Screenshot 2024-03-13 at 6 09 36 PM"
src="https://github.com/pinecone-io/pinecone-java-client/assets/119623786/90f46e4a-4dfd-4638-9c17-4f6bfa9b9432">

Next steps would be to create a Junit extension and possible
`IndexManagerSingleton` to manage index resources across all tests
directly. This would help make things more predictable and reliable when
adding future tests. This would also allow us to handle concurrent
integration test runs under the same API key, which is currently very
difficult due to `findIndexByDimensionAndType`.

## Type of Change
- [X] Bug fix (non-breaking change which fixes an issue)
- [X] Infrastructure change (CI configs, etc)

## Test Plan
Run integration tests and compare over run length and

Author: austin-denoble

.github/workflows/pr.yml

@@ -80,7 +80,7 @@ jobs:
       run: gradle clean build
 
     - name: Run integration tests
-      run: gradle integrationTest
+      run: gradle integrationTest --info
       env:
         PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }}
         PINECONE_ENVIRONMENT: ${{ secrets.PINECONE_ENVIRONMENT }}

src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java

@@ -6,6 +6,7 @@
 import io.pinecone.proto.SparseValues;
 import io.pinecone.proto.UpsertRequest;
 import io.pinecone.proto.Vector;
+import io.pinecone.unsigned_indices_model.VectorWithUnsignedIndices;
 
 import java.util.*;
 
@@ -135,21 +136,15 @@ public static UpsertRequest buildRequiredUpsertRequest(List<String> upsertIds, S
                 .build();
     }
 
-    public static UpsertRequest buildRequiredUpsertRequestByDimension(List<String> upsertIds, int dimension, String namespace) {
+    public static List<VectorWithUnsignedIndices> buildRequiredUpsertRequestByDimension(List<String> upsertIds, int dimension) {
         if (upsertIds.isEmpty()) upsertIds = Arrays.asList("v1", "v2", "v3");
 
-        List<Vector> upsertVectors = new ArrayList<>();
+        List<VectorWithUnsignedIndices> upsertVectors = new ArrayList<>();
         for (String upsertId : upsertIds) {
-            upsertVectors.add(Vector.newBuilder()
-                    .addAllValues(generateVectorValuesByDimension(dimension))
-                    .setId(upsertId)
-                    .build());
+            upsertVectors.add(new VectorWithUnsignedIndices(upsertId, generateVectorValuesByDimension(dimension)));
         }
 
-        return UpsertRequest.newBuilder()
-                .addAllVectors(upsertVectors)
-                .setNamespace(namespace)
-                .build();
+        return upsertVectors;
     }
 
     public static UpsertRequest buildOptionalUpsertRequest() {

src/integration/java/io/pinecone/helpers/IndexManager.java

@@ -2,70 +2,50 @@
 
 import io.pinecone.clients.Index;
 import io.pinecone.clients.Pinecone;
-import io.pinecone.configs.PineconeConfig;
-import io.pinecone.configs.PineconeConnection;
 import io.pinecone.exceptions.PineconeException;
 import org.openapitools.client.model.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.AbstractMap;
 import java.util.List;
 
-import static io.pinecone.helpers.AssertRetry.assertWithRetry;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class IndexManager {
     private static final Logger logger = LoggerFactory.getLogger(IndexManager.class);
 
-    public static PineconeConnection createIndexIfNotExistsDataPlane(int dimension, String indexType) throws IOException, InterruptedException {
+    public static AbstractMap.SimpleEntry<String, Pinecone> createIndexIfNotExistsDataPlane(int dimension, String indexType) throws IOException, InterruptedException {
         String apiKey = System.getenv("PINECONE_API_KEY");
         Pinecone pinecone = new Pinecone(apiKey);
 
         String indexName = findIndexWithDimensionAndType(pinecone, dimension, indexType);
-        if (indexName.isEmpty()) indexName = createNewIndex(pinecone, dimension, indexType);
+        if (indexName.isEmpty()) indexName = createNewIndex(pinecone, dimension, indexType, true);
 
-        // Do not proceed until the newly created index is ready
-        isIndexReady(indexName, pinecone);
-
-        // Adding to test PineconeConnection(pineconeConfig, host) constructor
-        String host = pinecone.describeIndex(indexName).getHost();
-        PineconeConfig config = new PineconeConfig(apiKey);
-        config.setHost(host);
-        return new PineconeConnection(config);
-    }
-
-    public static String createIndexIfNotExistsControlPlane(Pinecone pinecone, int dimension, String indexType) throws IOException, InterruptedException {
-        String indexName = findIndexWithDimensionAndType(pinecone, dimension, indexType);
-
-        return (indexName.isEmpty()) ? createNewIndex(pinecone, dimension, indexType) : indexName;
+        return new AbstractMap.SimpleEntry<>(indexName, pinecone);
     }
 
-    public static String findIndexWithDimensionAndType(Pinecone pinecone, int dimension, String indexType)
-            throws InterruptedException {
+    public static String findIndexWithDimensionAndType(Pinecone pinecone, int dimension, String indexType) {
         String indexName = "";
-        int i = 0;
         List<IndexModel> indexModels = pinecone.listIndexes().getIndexes();
         if(indexModels == null) {
             return indexName;
         }
-        while (i < indexModels.size()) {
-             IndexModel indexModel = isIndexReady(indexModels.get(i).getName(), pinecone);
+
+        for (IndexModel indexModel : indexModels) {
             if (indexModel.getDimension() == dimension
-                    && ((indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_POD)
-                        && indexModel.getSpec().getPod() != null
-                        && indexModel.getSpec().getPod().getReplicas() == 1
-                        && indexModel.getSpec().getPod().getPodType().equalsIgnoreCase("p1.x1"))
-                    || (indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_SERVERLESS)))) {
+                && (indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_POD) && indexModel.getSpec().getPod() != null)
+                || (indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_SERVERLESS) && indexModel.getSpec().getServerless() != null)
+            ) {
                 return indexModel.getName();
             }
-            i++;
         }
         return indexName;
     }
 
-    public static String createNewIndex(Pinecone pinecone, int dimension, String indexType) {
+    public static String createNewIndex(Pinecone pinecone, int dimension, String indexType, boolean waitUntilIndexIsReady) throws InterruptedException {
         String indexName = RandomStringBuilder.build("index-name", 8);
         String environment = System.getenv("PINECONE_ENVIRONMENT");
         CreateIndexRequestSpec createIndexRequestSpec;
@@ -74,7 +54,9 @@ public static String createNewIndex(Pinecone pinecone, int dimension, String ind
             CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).podType("p1.x1");
             createIndexRequestSpec = new CreateIndexRequestSpec().pod(podSpec);
         } else {
-            ServerlessSpec serverlessSpec = new ServerlessSpec().cloud(ServerlessSpec.CloudEnum.AWS).region(environment);
+            // Serverless currently has limited availability in specific regions, hardcode us-west-2 for now
+            ServerlessSpec serverlessSpec =
+                    new ServerlessSpec().cloud(ServerlessSpec.CloudEnum.AWS).region("us-west-2");
             createIndexRequestSpec = new CreateIndexRequestSpec().serverless(serverlessSpec);
         }
 
@@ -85,46 +67,61 @@ public static String createNewIndex(Pinecone pinecone, int dimension, String ind
                 .spec(createIndexRequestSpec);
         pinecone.createIndex(createIndexRequest);
 
+        if (waitUntilIndexIsReady) {
+            waitUntilIndexIsReady(pinecone, indexName);
+        }
+
         return indexName;
     }
 
     public static IndexModel waitUntilIndexIsReady(Pinecone pinecone, String indexName, Integer totalMsToWait) throws InterruptedException {
         IndexModel index = pinecone.describeIndex(indexName);
         int waitedTimeMs = 0;
-        int intervalMs = 1500;
+        int intervalMs = 2000;
 
         while (!index.getStatus().getReady()) {
             index = pinecone.describeIndex(indexName);
             if (waitedTimeMs >= totalMsToWait) {
-                logger.info("Index " + indexName + " not ready after " + waitedTimeMs + "ms");
+                logger.info("WARNING: Index " + indexName + " not ready after " + waitedTimeMs + "ms");
                 break;
             }
             if (index.getStatus().getReady()) {
                 logger.info("Index " + indexName + " is ready after " + waitedTimeMs + "ms");
+                // Wait one final time before we start connecting and operating on the index
+                Thread.sleep(10000);
                 break;
             }
             Thread.sleep(intervalMs);
+            logger.info("Waited " + waitedTimeMs + "ms for " + indexName + " to get ready");
             waitedTimeMs += intervalMs;
         }
         return index;
     }
 
     public static IndexModel waitUntilIndexIsReady(Pinecone pinecone, String indexName) throws InterruptedException {
-        return waitUntilIndexIsReady(pinecone, indexName, 120000);
+        return waitUntilIndexIsReady(pinecone, indexName, 200000);
     }
 
-    public static PineconeConnection createNewIndexAndConnect(Pinecone pinecone, String indexName, int dimension, IndexMetric metric, CreateIndexRequestSpec spec) throws InterruptedException, PineconeException {
-        String apiKey = System.getenv("PINECONE_API_KEY");
+    public static Pinecone createNewIndex(Pinecone pinecone, String indexName, int dimension, IndexMetric metric, CreateIndexRequestSpec spec, boolean waitUntilIndexIsReady) throws InterruptedException, PineconeException {
+        CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(indexName).dimension(dimension).metric(metric).spec(spec);
+        pinecone.createIndex(createIndexRequest);
+
+        if (waitUntilIndexIsReady) {
+            waitUntilIndexIsReady(pinecone, indexName);
+        }
+        return pinecone;
+    }
+
+    public static Index createNewIndexAndConnect(Pinecone pinecone, String indexName, int dimension, IndexMetric metric, CreateIndexRequestSpec spec) throws InterruptedException, PineconeException {
         CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(indexName).dimension(dimension).metric(metric).spec(spec);
         pinecone.createIndex(createIndexRequest);
 
         // Wait until index is ready
-        waitUntilIndexIsReady(pinecone, indexName, 200000);
+        waitUntilIndexIsReady(pinecone, indexName);
         // wait a bit more before we connect...
-        Thread.sleep(15000);
+        Thread.sleep(5000);
 
-        PineconeConfig config = new PineconeConfig(apiKey);
-        return new PineconeConnection(config, indexName);
+        return pinecone.createIndexConnection(indexName);
     }
 
     public static CollectionModel createCollection(Pinecone pinecone, String collectionName, String indexName, boolean waitUntilReady) throws InterruptedException {
@@ -138,7 +135,8 @@ public static CollectionModel createCollection(Pinecone pinecone, String collect
             int timeWaited = 0;
             CollectionModel.StatusEnum collectionReady = collection.getStatus();
             while (collectionReady != CollectionModel.StatusEnum.READY && timeWaited < 120000) {
-                logger.info("Waiting for collection " + collectionName + " to be ready. Waited " + timeWaited + " milliseconds...");
+                logger.info("Waiting for collection " + collectionName + " to be ready. Waited " + timeWaited + " " +
+                        "milliseconds...");
                 Thread.sleep(5000);
                 timeWaited += 5000;
                 collection = pinecone.describeCollection(collectionName);
@@ -152,15 +150,4 @@ public static CollectionModel createCollection(Pinecone pinecone, String collect
 
         return collection;
     }
-
-    public static IndexModel isIndexReady(String indexName, Pinecone pinecone)
-            throws InterruptedException {
-        final IndexModel[] indexModels = new IndexModel[1];
-        assertWithRetry(() -> {
-            indexModels[0] = pinecone.describeIndex(indexName);
-            assert (indexModels[0].getStatus().getReady());
-        }, 4);
-
-        return indexModels[0];
-    }
 }