Skip to content

Commit

Permalink
Add describeIndexStats integration test (#47)
Browse files Browse the repository at this point in the history
## Problem

The java sdk lacks integration test for describeIndexStats functionality

## Solution

Added integration test for both future and sync client to
describeIndexStats. Note this test does include Upsert operation but
I'll be testing it rigorously in UpsertAndFetch integration 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)
- [X] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan

Ran integration tests


Note: PineconeClientLiveIntegTest will go away once I add all of the
integration tests for data plane.
PineconeIndexOperationsClientIntegrationTest will be moved under control
plane integration test sub-folder and renamed to CreateAndDeleteIndex,
but I'll be adding more tests in their if needed.
  • Loading branch information
rohanshah18 authored Nov 24, 2023
1 parent 792cb4f commit d1546e1
Show file tree
Hide file tree
Showing 4 changed files with 259 additions and 43 deletions.
60 changes: 60 additions & 0 deletions src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
package io.pinecone.helpers;

import com.google.common.primitives.Floats;
import com.google.protobuf.Struct;
import com.google.protobuf.Value;
import io.pinecone.proto.SparseValues;
import io.pinecone.proto.UpsertRequest;
import io.pinecone.proto.Vector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class BuildUpsertRequest {
private static final float[][] upsertData = {{1.0F, 2.0F, 3.0F}, {4.0F, 5.0F, 6.0F}, {7.0F, 8.0F, 9.0F}};

public static UpsertRequest buildRequiredUpsertRequest() {
String namespace = RandomStringBuilder.build("ns", 8);
List<String> upsertIds = Arrays.asList("v1", "v2", "v3");
List<Vector> upsertVectors = new ArrayList<>();

for (int i = 0; i < upsertData.length; i++) {
upsertVectors.add(Vector.newBuilder()
.addAllValues(Floats.asList(upsertData[i]))
.setMetadata(Struct.newBuilder()
.putFields("some_field", Value.newBuilder().setNumberValue(i).build())
.build())
.setId(upsertIds.get(i))
.build());
}

return UpsertRequest.newBuilder()
.addAllVectors(upsertVectors)
.setNamespace(namespace)
.build();
}

public static UpsertRequest buildOptionalUpsertRequest() {
String namespace = RandomStringBuilder.build("ns", 8);
List<String> hybridsIds = Arrays.asList("v4", "v5", "v6");
List<Vector> hybridVectors = new ArrayList<>();
List<Integer> sparseIndices = Arrays.asList(0, 1, 2);
List<Float> sparseValues = Arrays.asList(0.11f, 0.22f, 0.33f);
for (int i = 0; i < hybridsIds.size(); i++) {
hybridVectors.add(
Vector.newBuilder()
.addAllValues(Floats.asList(upsertData[i]))
.setSparseValues(
SparseValues.newBuilder().addAllIndices(sparseIndices).addAllValues(sparseValues).build()
)
.setId(hybridsIds.get(i))
.build());
}

return UpsertRequest.newBuilder()
.addAllVectors(hybridVectors)
.setNamespace(namespace)
.build();
}
}
59 changes: 59 additions & 0 deletions src/integration/java/io/pinecone/helpers/IndexManager.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
package io.pinecone.helpers;

import io.pinecone.*;
import io.pinecone.model.CreateIndexRequest;
import io.pinecone.model.IndexMeta;

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

public class IndexManager {
public PineconeConnection createIndexIfNotExists(int dimension) throws IOException, InterruptedException {
boolean createNewIndex = false;
String indexName = "";
PineconeClientConfig config = new PineconeClientConfig()
.withApiKey(System.getenv("PINECONE_API_KEY"))
.withEnvironment(System.getenv("PINECONE_ENVIRONMENT"));
PineconeIndexOperationClient controlPlaneClient = new PineconeIndexOperationClient(config);
List<String> 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);
}

PineconeClient dataPlaneClient = new PineconeClient(config);
IndexMeta indexMeta = controlPlaneClient.describeIndex(indexName);
String host = indexMeta.getStatus().getHost();

return dataPlaneClient.connect(
new PineconeConnectionConfig()
.withConnectionUrl("https://" + host));
}

public static IndexMeta isIndexReady(String indexName, PineconeIndexOperationClient indexOperationClient)
throws IOException, InterruptedException {
IndexMeta indexMeta;
while (true) {
indexMeta = indexOperationClient.describeIndex(indexName);
if (indexMeta.getStatus().getState().equalsIgnoreCase("ready")) {
break;
}

Thread.sleep(1000);
}
return indexMeta;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,11 @@

import java.io.IOException;

import static io.pinecone.helpers.IndexManager.isIndexReady;
import static org.junit.jupiter.api.Assertions.assertEquals;

public class ConfigureIndexTest {
private PineconeIndexOperationClient pinecone;
private PineconeIndexOperationClient indexOperationClient;
private String indexName;
private static final Logger logger = LoggerFactory.getLogger(PineconeClientLiveIntegTest.class);

Expand All @@ -28,28 +29,29 @@ public void setUp() throws IOException {
PineconeClientConfig config = new PineconeClientConfig()
.withApiKey(System.getenv("PINECONE_API_KEY"))
.withEnvironment(System.getenv("PINECONE_ENVIRONMENT"));
pinecone = new PineconeIndexOperationClient(config);
indexOperationClient = new PineconeIndexOperationClient(config);

// Create an index
CreateIndexRequest request = new CreateIndexRequest()
.withIndexName(indexName)
.withDimension(5)
.withMetric("euclidean");
pinecone.createIndex(request);
indexOperationClient.createIndex(request);
}

@AfterEach
public void cleanUp() throws IOException {
pinecone.deleteIndex(indexName);
public void cleanUp() throws IOException, InterruptedException {
indexOperationClient.deleteIndex(indexName);
Thread.sleep(3500);
}

@Test
public void configureIndexWithInvalidIndexName() {
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withReplicas(2);
try {
checkIndexStatusReady(indexName);
pinecone.configureIndex("non-existent-index", configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex("non-existent-index", configureIndexRequest);
} catch (PineconeNotFoundException exception) {
assert (exception.getLocalizedMessage().contains("404: Not Found"));
} catch (IOException | InterruptedException exception) {
Expand All @@ -62,8 +64,8 @@ public void configureIndexExceedingQuota() {
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withReplicas(20);
try {
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);
} catch (PineconeBadRequestException exception) {
assert (exception.getLocalizedMessage().contains("The index exceeds the project quota"));
assert (exception.getLocalizedMessage().contains("Upgrade your account or change" +
Expand All @@ -77,17 +79,17 @@ public void configureIndexExceedingQuota() {
public void scaleUp() {
try{
// Verify the starting state
IndexMeta indexMeta = pinecone.describeIndex(indexName);
IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
assertEquals(1, indexMeta.getDatabase().getReplicas());

// Configure the index
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withReplicas(2);
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);

// Verify replicas were scaled up
indexMeta = pinecone.describeIndex(indexName);
indexMeta = indexOperationClient.describeIndex(indexName);
assertEquals(2, indexMeta.getDatabase().getReplicas());
} catch (Exception exception) {
logger.error(exception.toString());
Expand All @@ -98,27 +100,27 @@ public void scaleUp() {
public void scaleDown() {
try {
// Verify the starting state
IndexMeta indexMeta = pinecone.describeIndex(indexName);
IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
assertEquals(1, indexMeta.getDatabase().getReplicas());

// Scale up for the test
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withReplicas(3);
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);

// Verify the scaled up replicas
indexMeta = pinecone.describeIndex(indexName);
indexMeta = indexOperationClient.describeIndex(indexName);
assertEquals(3, indexMeta.getDatabase().getReplicas());

// Scaling down
configureIndexRequest = new ConfigureIndexRequest()
.withReplicas(1);
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);

// Verify replicas were scaled down
indexMeta = pinecone.describeIndex(indexName);
indexMeta = indexOperationClient.describeIndex(indexName);
assertEquals(1, indexMeta.getDatabase().getReplicas());
} catch (Exception exception) {
logger.error(exception.toString());
Expand All @@ -129,15 +131,15 @@ public void scaleDown() {
public void changingBasePodType() {
try {
// Verify the starting state
IndexMeta indexMeta = pinecone.describeIndex(indexName);
IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
assertEquals("p1.x1", indexMeta.getDatabase().getPodType());

// Try to change the base pod type
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withPodType("p2.x1");

checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);
} catch (PineconeBadRequestException pineconeBadRequestException) {
assertEquals(pineconeBadRequestException.getMessage(), "updating base pod type is not supported");
} catch (Exception exception) {
Expand All @@ -149,17 +151,17 @@ public void changingBasePodType() {
public void sizeIncrease() {
try {
// Verify the starting state
IndexMeta indexMeta = pinecone.describeIndex(indexName);
IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
assertEquals("p1.x1", indexMeta.getDatabase().getPodType());

// Change the pod type to a larger one
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withPodType("p1.x2");
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);

// Get the index description to verify the new pod type
indexMeta = pinecone.describeIndex(indexName);
indexMeta = indexOperationClient.describeIndex(indexName);
assertEquals("p1.x2", indexMeta.getDatabase().getPodType());
} catch (Exception exception) {
logger.error(exception.getLocalizedMessage());
Expand All @@ -170,37 +172,26 @@ public void sizeIncrease() {
public void sizeDown() {
try {
// Verify the starting state
IndexMeta indexMeta = pinecone.describeIndex(indexName);
IndexMeta indexMeta = isIndexReady(indexName, indexOperationClient);
assertEquals("p1.x1", indexMeta.getDatabase().getPodType());

// Increase the pod type
ConfigureIndexRequest configureIndexRequest = new ConfigureIndexRequest()
.withPodType("p1.x2");
checkIndexStatusReady(indexName);
pinecone.configureIndex(indexName, configureIndexRequest);
isIndexReady(indexName, indexOperationClient);
indexOperationClient.configureIndex(indexName, configureIndexRequest);

// Get the index description to verify the new pod type
indexMeta = pinecone.describeIndex(indexName);
indexMeta = indexOperationClient.describeIndex(indexName);
assertEquals("p1.x2", indexMeta.getDatabase().getPodType());

// Attempt to scale down
configureIndexRequest = new ConfigureIndexRequest()
.withPodType("p1.x1");
pinecone.configureIndex(indexName, configureIndexRequest);
indexOperationClient.configureIndex(indexName, configureIndexRequest);
} catch (Exception exception) {
assertEquals(exception.getClass(), PineconeBadRequestException.class);
assertEquals(exception.getMessage(), "scaling down pod type is not supported");
}
}

private void checkIndexStatusReady(String indexName) throws IOException, InterruptedException {
while (true) {
IndexMeta indexMeta = pinecone.describeIndex(indexName);
if (indexMeta.getStatus().getState().equalsIgnoreCase("ready")) {
break;
}

Thread.sleep(1000);
}
}
}
Loading

0 comments on commit d1546e1

Please sign in to comment.