handle future listenables properly in error tests

Author: austin-denoble

src/integration/java/io/pinecone/integration/dataPlane/QueryErrorPodTest.java

@@ -1,12 +1,14 @@
 package io.pinecone.integration.dataPlane;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import io.grpc.StatusRuntimeException;
 import io.pinecone.clients.Index;
 import io.pinecone.clients.Pinecone;
 import io.pinecone.clients.AsyncIndex;
 import io.pinecone.exceptions.PineconeValidationException;
 import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.proto.*;
+import io.pinecone.unsigned_indices_model.QueryResponseWithUnsignedIndices;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.openapitools.client.model.IndexModelSpec;
@@ -16,6 +18,8 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static io.pinecone.helpers.BuildUpsertRequest.*;
 import static io.pinecone.helpers.IndexManager.createIndexIfNotExistsDataPlane;
@@ -76,15 +80,16 @@ public void QueryWithNullSparseIndicesNotNullSparseValuesSyncTest() {
     }
 
     @Test
-    public void queryWithIncorrectVectorDimensionFuture() throws ExecutionException, InterruptedException {
+    public void queryWithIncorrectVectorDimensionFuture() throws ExecutionException, InterruptedException, TimeoutException {
         String namespace = RandomStringBuilder.build("ns", 8);
         DescribeIndexStatsResponse describeIndexStatsResponse1 = asyncIndexClient.describeIndexStats(null).get();
         assertEquals(describeIndexStatsResponse1.getDimension(), dimension);
 
         // Query with incorrect dimensions
         try {
             List<Float> vector = Arrays.asList(100F);
-            asyncIndexClient.query(5, vector, null, null, null, namespace, null, true, true).get();
+            ListenableFuture<QueryResponseWithUnsignedIndices> queryFuture = asyncIndexClient.query(5, vector, null, null, null, namespace, null, true, true);
+            queryFuture.get(10, TimeUnit.SECONDS);
 
             fail("queryWithIncorrectVectorDimensionFuture should have thrown ExecutionException");
         } catch (ExecutionException expected) {
@@ -94,20 +99,22 @@ public void queryWithIncorrectVectorDimensionFuture() throws ExecutionException,
     }
 
     @Test
-    public void QueryWithNullSparseIndicesNotNullSparseValuesFutureTest() throws ExecutionException, InterruptedException {
+    public void QueryWithNullSparseIndicesNotNullSparseValuesFutureTest() throws ExecutionException, InterruptedException, TimeoutException {
         String id = RandomStringBuilder.build(3);
 
         try {
-            asyncIndexClient.update(id,
+            ListenableFuture<UpdateResponse> updateFuture = asyncIndexClient.update(id,
                     generateVectorValuesByDimension(dimension),
                     null,
                     null,
                     null,
-                    generateVectorValuesByDimension(dimension)).get();
+                    generateVectorValuesByDimension(dimension));
+
+            updateFuture.get(10, TimeUnit.SECONDS);
 
             fail("QueryWithNullSparseIndicesNotNullSparseValuesFutureTest should have thrown PineconeValidationException");
         } catch (PineconeValidationException expected) {
-            assertEquals(expected.getLocalizedMessage(), "ensure that both sparse indices and values are present");
+            assertTrue(expected.getLocalizedMessage().contains("ensure that both sparse indices and values are present"));
         }
     }
 }

src/integration/java/io/pinecone/integration/dataPlane/UpdateFetchAndQueryPodTest.java

@@ -1,5 +1,6 @@
 package io.pinecone.integration.dataPlane;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.protobuf.Struct;
 import com.google.protobuf.Value;
 import io.grpc.StatusRuntimeException;
@@ -21,6 +22,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static io.pinecone.helpers.AssertRetry.assertWithRetry;
 import static io.pinecone.helpers.BuildUpsertRequest.*;
@@ -259,7 +262,7 @@ public void updateNullSparseIndicesNotNullSparseValuesSyncTest() {
 
             fail("updateNullSparseIndicesNotNullSparseValuesSyncTest should have thrown PineconeValidationException");
         } catch (PineconeValidationException expected) {
-            assertEquals(expected.getLocalizedMessage(), "ensure that both sparse indices and values are present");
+            assertTrue(expected.getLocalizedMessage().contains("ensure that both sparse indices and values are present"));
         }
     }
 
@@ -385,7 +388,7 @@ public void updateAllParamsFetchAndQueryFutureTest() throws InterruptedException
     }
 
     @Test
-    public void addIncorrectDimensionalValuesFutureTest() throws InterruptedException {
+    public void addIncorrectDimensionalValuesFutureTest() throws ExecutionException, InterruptedException, TimeoutException {
         // Upsert vectors with required parameters
         int numOfVectors = 3;
         String namespace = RandomStringBuilder.build("ns", 8);
@@ -412,7 +415,8 @@ public void addIncorrectDimensionalValuesFutureTest() throws InterruptedExceptio
 
         // Should fail since only 1 value is added for the vector of dimension 3
         try {
-            asyncIndexClient.update(idToUpdate, updatedValues, null, namespace, null, null);
+            ListenableFuture<UpdateResponse> updateFuture = asyncIndexClient.update(idToUpdate, updatedValues, null, namespace, null, null);
+            updateFuture.get(10, TimeUnit.SECONDS);
 
             fail("addIncorrectDimensionalValuesFutureTest should have thrown StatusRuntimeException");
         } catch (StatusRuntimeException expected) {
@@ -466,20 +470,22 @@ public void queryWithFilersFutureTest() throws ExecutionException, InterruptedEx
     }
 
     @Test
-    public void updateNullSparseIndicesNotNullSparseValuesFutureTest() {
+    public void updateNullSparseIndicesNotNullSparseValuesFutureTest() throws ExecutionException, InterruptedException, TimeoutException {
         String id = RandomStringBuilder.build(3);
 
         try {
-            asyncIndexClient.update(id,
+            ListenableFuture<UpdateResponse> updateFuture = asyncIndexClient.update(id,
                     generateVectorValuesByDimension(dimension),
                     null,
                     null,
                     null,
                     generateVectorValuesByDimension(dimension));
+            updateFuture.get(10, TimeUnit.SECONDS);
+
 
             fail("updateNullSparseIndicesNotNullSparseValuesFutureTest should have thrown PineconeValidationException");
         } catch (PineconeValidationException expected) {
-            assertEquals(expected.getLocalizedMessage(), "ensure that both sparse indices and values are present");
+            assertTrue(expected.getLocalizedMessage().contains("ensure that both sparse indices and values are present"));
         }
     }
 }