Skip to content

Commit

Permalink
fix issue with assertWithRetry incrementation, add new waitUntilIndex…
Browse files Browse the repository at this point in the history 
…Ready, clean up helper functions and tests
  • Loading branch information
@austin-denoble
austin-denoble committed Feb 12, 2024
1 parent 9a6abcc commit 6523554
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 24 deletions.
9 changes: 5 additions & 4 deletions src/integration/java/io/pinecone/helpers/AssertRetry.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,16 @@
import java.util.concurrent.ExecutionException;

public class AssertRetry {
private static final int maxRetry = 5;
private static int delay = 1500;
private static final int maxRetry = 4;
private static final int delay = 1500;

public static void assertWithRetry(AssertionRunnable assertionRunnable) throws InterruptedException, PineconeException {
assertWithRetry(assertionRunnable, 2);
}

public static void assertWithRetry(AssertionRunnable assertionRunnable, int backOff) throws InterruptedException, PineconeException {
int retryCount = 0;
int delayCount = delay;
boolean success = false;

while (retryCount < maxRetry && !success) {
Expand All @@ -24,8 +25,8 @@ public static void assertWithRetry(AssertionRunnable assertionRunnable, int back
success = true;
} catch (AssertionError | ExecutionException | IOException e) {
retryCount++;
delay*=backOff;
Thread.sleep(delay);
Thread.sleep(delayCount);
delayCount = delayCount * backOff;
}
}
}
Expand Down
28 changes: 27 additions & 1 deletion src/integration/java/io/pinecone/helpers/IndexManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,11 +109,37 @@ public static IndexModel isIndexReady(String indexName, PineconeControlPlaneClie
assertWithRetry(() -> {
indexModels[0] = controlPlaneClient.describeIndex(indexName);
assert (indexModels[0].getStatus().getReady());
}, 2);
}, 3);

return indexModels[0];
}

public static IndexModel waitUntilIndexIsReady(PineconeControlPlaneClient controlPlaneClient, String indexName, Integer totalMsToWait) throws InterruptedException {
IndexModel index = controlPlaneClient.describeIndex(indexName);
int waitedTimeMs = 0;
int intervalMs = 1500;

while (!index.getStatus().getReady()) {
index = controlPlaneClient.describeIndex(indexName);
if (waitedTimeMs >= totalMsToWait) {
System.out.println("Index " + indexName + " not ready after " + waitedTimeMs + "ms");
break;
}
if (index.getStatus().getReady()) {
Thread.sleep(2500);
System.out.println("Index " + indexName + " is ready after " + waitedTimeMs + "ms");
break;
}
Thread.sleep(intervalMs);
waitedTimeMs += intervalMs;
}
return index;
}

public static IndexModel waitUntilIndexIsReady(PineconeControlPlaneClient controlPlaneClient, String indexName) throws InterruptedException {
return waitUntilIndexIsReady(controlPlaneClient, indexName, 120000);
}

public static CollectionModel createCollection(PineconeControlPlaneClient controlPlaneClient, String collectionName, String indexName, boolean waitUntilReady) throws InterruptedException {
CreateCollectionRequest createCollectionRequest = new CreateCollectionRequest().name(collectionName).source(indexName);
CollectionModel collection = controlPlaneClient.createCollection(createCollectionRequest);
Expand Down
10 changes: 4 additions & 6 deletions src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,10 @@
public class CollectionErrorTest {
private static final String apiKey = System.getenv("PINECONE_API_KEY");
private static final String environment = System.getenv("PINECONE_ENVIRONMENT");
private static final String indexName = RandomStringBuilder.build("collection-error-test-", 8);
private static final String indexName = RandomStringBuilder.build("collection-error-test", 8);
private static final List<String> upsertIds = Arrays.asList("v1", "v2", "v3");
private static final int dimension = 3;
private static final String collectionName = RandomStringBuilder.build("reusable-coll-", 8);
private static final String collectionName = RandomStringBuilder.build("reusable-coll", 8);
private static CollectionModel collection;
private static PineconeControlPlaneClient controlPlaneClient;

Expand Down Expand Up @@ -60,7 +60,7 @@ public static void cleanUp() {
@Test
public void testCreateCollectionFromInvalidIndex() {
try {
CreateCollectionRequest createCollectionRequest = new CreateCollectionRequest().name(RandomStringBuilder.build("coll1-", 8)).source("invalid-index");
CreateCollectionRequest createCollectionRequest = new CreateCollectionRequest().name(RandomStringBuilder.build("coll1", 8)).source("invalid-index");
controlPlaneClient.createCollection(createCollectionRequest);
} catch (PineconeException exception) {
assertTrue(exception.getMessage().contains("Resource invalid-index not found"));
Expand All @@ -71,7 +71,7 @@ public void testIndexFromNonExistentCollection() {
try {
CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection("non-existent-collection");
CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec);
CreateIndexRequest newCreateIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-nonexistent-coll-", 8)).dimension(3).metric(IndexMetric.COSINE).spec(spec);
CreateIndexRequest newCreateIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-nonexistent-coll", 8)).dimension(3).metric(IndexMetric.COSINE).spec(spec);
controlPlaneClient.createIndex(newCreateIndexRequest);
} catch (PineconeException exception) {
assertTrue(exception.getMessage().contains("Resource non-existent-collection not found"));
Expand Down Expand Up @@ -116,7 +116,6 @@ public void testCreateIndexWithMismatchedDimension() {
CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll-", 8)).dimension(dimension + 1).metric(IndexMetric.COSINE).spec(spec);
controlPlaneClient.createIndex(createIndexRequest);
} catch (PineconeException exception) {
System.out.println("Exception: " + exception);
assertTrue(exception.getMessage().contains("Index and collection must have the same dimension"));
}
}
Expand All @@ -136,5 +135,4 @@ public void testCreateCollectionFromNotReadyIndex() {
assertTrue(exception.getMessage().contains("Source index is not ready"));
}
}

}
28 changes: 15 additions & 13 deletions src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,13 @@
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.openapitools.client.ApiException;
import org.openapitools.client.model.*;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;

import static io.pinecone.helpers.IndexManager.createNewIndexAndConnect;
import static io.pinecone.helpers.IndexManager.isIndexReady;
import static io.pinecone.helpers.IndexManager.waitUntilIndexIsReady;
import static io.pinecone.helpers.IndexManager.createCollection;
import static io.pinecone.helpers.BuildUpsertRequest.*;
import static org.junit.jupiter.api.Assertions.*;
Expand Down Expand Up @@ -55,7 +53,7 @@ public void testIndexToCollectionHappyPath() throws InterruptedException {
String collectionName = RandomStringBuilder.build("collection-test", 8);

// Create collection from index
CollectionModel collection = createCollection(controlPlaneClient, indexName, collectionName, true);
CollectionModel collection = createCollection(controlPlaneClient, collectionName, indexName, true);

assertEquals(collection.getName(), collectionName);
assertEquals(collection.getEnvironment(), environment);
Expand Down Expand Up @@ -87,18 +85,20 @@ public void testIndexToCollectionHappyPath() throws InterruptedException {
assertTrue(collection.getSize() > 0);

// Create index from collection
String newIndexName = "index-from-collection-" + collectionName;
String newIndexName = RandomStringBuilder.build("index-from-col-", 5);
System.out.println("Creating index " + newIndexName + " from collection " + collectionName);

CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection(collectionName);
CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec);
CreateIndexRequest newCreateIndexRequest = new CreateIndexRequest().name(newIndexName).dimension(dimension).metric(indexMetric).spec(spec);
IndexModel indexFromCollection = controlPlaneClient.createIndex(newCreateIndexRequest);
System.out.println("Index " + newIndexName + " created from collection " + collectionName + ". Waiting until index is ready...");
isIndexReady(newIndexName, controlPlaneClient);
indexFromCollection = waitUntilIndexIsReady(controlPlaneClient, newIndexName);
System.out.println("Index " + newIndexName + " is ready");

IndexModel indexDescription = controlPlaneClient.describeIndex(newIndexName);
assertEquals(indexDescription.getName(), newIndexName);
assertEquals(indexDescription.getSpec().getPod().getSourceCollection(), collectionName);
assertEquals(indexDescription.getStatus().getReady(), true);

// Set up new index data plane connection
Expand All @@ -107,20 +107,22 @@ public void testIndexToCollectionHappyPath() throws InterruptedException {
DescribeIndexStatsResponse describeResponse = newIndexDataPlaneClient.getBlockingStub().describeIndexStats(DescribeIndexStatsRequest.newBuilder().build());

// Verify stats reflect the vectors in the collection
System.out.println("Index " + newIndexName + " stats: " + describeResponse);
assertEquals(describeResponse.getTotalVectorCount(), 3);

// Verify the vectors from the collection -> new index can be fetched
FetchResponse fetchedVectors = newIndexDataPlaneClient.getBlockingStub().fetch(FetchRequest.newBuilder().addAllIds(upsertIds).setNamespace(namespace).build());
newIndexDataPlaneClient.close();

for (String key : upsertIds) {
assert (fetchedVectors.containsVectors(key));
}

// Verify we can delete the collection
controlPlaneClient.deleteCollection(collectionName);
Thread.sleep(2500);
collections = controlPlaneClient.listCollections().getCollections();


if (collections != null) {
boolean isCollectionDeleted = true;
for (CollectionModel col : collections) {
Expand All @@ -134,18 +136,17 @@ public void testIndexToCollectionHappyPath() throws InterruptedException {
fail("Collection " + collectionName + " was not successfully deleted");
}
}

// Clean up
controlPlaneClient.deleteIndex(newIndexName);
controlPlaneClient.deleteCollection(collectionName);
newIndexDataPlaneClient.close();
}

@Test
public void testIndexFromDifferentMetricCollection() throws InterruptedException {
String collectionName = RandomStringBuilder.build("collection-test", 8);

// Create collection from index
CollectionModel collection = createCollection(controlPlaneClient, indexName, collectionName, true);
CollectionModel collection = createCollection(controlPlaneClient, collectionName, indexName, true);

assertEquals(collection.getName(), collectionName);
assertEquals(collection.getEnvironment(), environment);
Expand All @@ -159,13 +160,14 @@ public void testIndexFromDifferentMetricCollection() throws InterruptedException
}
}

String indexName = RandomStringBuilder.build("from-coll-", 8);
String newIndexName = RandomStringBuilder.build("from-coll", 8);
CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection(collectionName);
CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec);
PineconeConnection dataPlaneConnection = createNewIndexAndConnect(controlPlaneClient, indexName, dimension, targetMetric, spec);
PineconeConnection dataPlaneConnection = createNewIndexAndConnect(controlPlaneClient, newIndexName, dimension, targetMetric, spec);
VectorServiceGrpc.VectorServiceBlockingStub blockingStub = dataPlaneConnection.getBlockingStub();

controlPlaneClient.deleteIndex(indexName);
controlPlaneClient.deleteIndex(newIndexName);
controlPlaneClient.deleteCollection(collectionName);
dataPlaneConnection.close();
}

Expand Down

0 comments on commit 6523554

Please sign in to comment.