make sure we're closing PineconeConnection objects properly in tests …

…where we create one
pinecone-io · Feb 14, 2024 · 45b8af2 · 45b8af2
1 parent 1bfd814
commit 45b8af2
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 7 deletions.
diff --git a/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java
@@ -46,7 +46,7 @@ public static void setUpIndexAndCollection() throws InterruptedException {
 
         // Upsert vectors to index and sleep for freshness
         blockingStub.upsert(buildRequiredUpsertRequestByDimension(upsertIds, dimension, ""));
-
+        dataPlaneConnection.close();
         Thread.sleep(3500);
 
         // Create collection from index
@@ -108,7 +108,7 @@ public void testCreateIndexInMismatchedEnvironment() {
 
             CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().sourceCollection(collection.getName()).environment(mismatchedEnv);
             CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec);
-            CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll-", 8)).dimension(dimension).metric(IndexMetric.COSINE).spec(spec);
+            CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll", 8)).dimension(dimension).metric(IndexMetric.COSINE).spec(spec);
             controlPlaneClient.createIndex(createIndexRequest);
         } catch (PineconeException exception) {
             assertTrue(exception.getMessage().contains("Source collection must be in the same environment as the index"));
@@ -122,7 +122,7 @@ public void testCreateIndexWithMismatchedDimension() {
             CollectionModel collection = controlPlaneClient.describeCollection(collectionName);
             CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().sourceCollection(collection.getName()).environment(collection.getEnvironment());
             CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec);
-            CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll-", 8)).dimension(dimension + 1).metric(IndexMetric.COSINE).spec(spec);
+            CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll", 8)).dimension(dimension + 1).metric(IndexMetric.COSINE).spec(spec);
             controlPlaneClient.createIndex(createIndexRequest);
         } catch (PineconeException exception) {
             assertTrue(exception.getMessage().contains("Index and collection must have the same dimension"));

diff --git a/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java
@@ -42,7 +42,9 @@ public static void setUp() throws InterruptedException {
 
         // Upsert vectors to index and sleep for freshness
         blockingStub.upsert(buildRequiredUpsertRequestByDimension(upsertIds, dimension, namespace));
+        dataPlaneConnection.close();
         Thread.sleep(3500);
+
     }
 
     @AfterAll

diff --git a/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java b/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java
@@ -6,6 +6,7 @@
 import io.pinecone.PineconeConnection;
 import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.proto.*;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -24,14 +25,21 @@
 public class PineconeClientLiveIntegTest {
 
     private static final Logger logger = LoggerFactory.getLogger(PineconeClientLiveIntegTest.class);
+
+    private static PineconeConnection connection;
     private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub;
 
     @BeforeAll
     public static void defineConfig() throws IOException, InterruptedException {
-        PineconeConnection connection = createIndexIfNotExistsDataPlane(3, IndexModelSpec.SERIALIZED_NAME_POD);
+        connection = createIndexIfNotExistsDataPlane(3, IndexModelSpec.SERIALIZED_NAME_POD);
         blockingStub = connection.getBlockingStub();
     }
 
+    @AfterAll
+    public static void cleanUp() {
+        connection.close();
+    }
+
     @Test
     public void sanity() throws InterruptedException {
         String namespace = RandomStringBuilder.build("ns", 8);

diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java
@@ -6,6 +6,8 @@
 import io.pinecone.PineconeConnection;
 import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.proto.*;
+import org.junit.After;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.openapitools.client.model.IndexModelSpec;
@@ -21,17 +23,23 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class UpdateAndQueryTest {
+    private static PineconeConnection connection;
     private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub;
     private static VectorServiceGrpc.VectorServiceFutureStub futureStub;
     private static final int dimension = 3;
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
+        connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }
 
+    @AfterAll
+    public static void cleanUp() {
+        connection.close();
+    }
+
     @Test
     public void UpdateRequiredParamsFetchAndQuerySync() throws InterruptedException {
         // Upsert vectors with required parameters

diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java
@@ -5,6 +5,7 @@
 import io.pinecone.PineconeConnection;
 import io.pinecone.helpers.RandomStringBuilder;
 import io.pinecone.proto.*;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.openapitools.client.model.IndexModelSpec;
@@ -20,17 +21,23 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class UpsertAndDeleteTest {
+    private static PineconeConnection connection;
     private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub;
     private static VectorServiceGrpc.VectorServiceFutureStub futureStub;
     private static final int dimension = 3;
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
+        connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }
 
+    @AfterAll
+    public static void cleanUp() {
+        connection.close();
+    }
+
     @Test
     public void UpsertVectorsAndDeleteByIdSyncTest() throws InterruptedException {
         // Upsert vectors with required parameters

diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java
@@ -14,17 +14,23 @@
 import java.util.concurrent.ExecutionException;
 
 public class UpsertAndDescribeIndexStatsTest {
+    private static PineconeConnection connection;
     private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub;
     private static VectorServiceGrpc.VectorServiceFutureStub futureStub;
     private static final int dimension = 3;
 
     @BeforeAll
     public static void setUp() throws IOException, InterruptedException {
-        PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
+        connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD);
         blockingStub = connection.getBlockingStub();
         futureStub = connection.getFutureStub();
     }
 
+    @AfterAll
+    public static void cleanUp() {
+        connection.close();
+    }
+
     @Test
     public void UpsertRequiredVectorsAndDescribeIndexStatsSyncTest() throws InterruptedException {
         // Get vector and namespace counts before upserting vectors with required parameters