Updated to address comments

googleapis · Dec 15, 2020 · 6ceee1c · 6ceee1c
1 parent e406335
commit 6ceee1c
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 4 deletions.
diff --git a/src/main/java/com/google/cloud/pubsublite/kafka/SharedBehavior.java b/src/main/java/com/google/cloud/pubsublite/kafka/SharedBehavior.java
@@ -28,8 +28,8 @@
 import org.apache.kafka.common.PartitionInfo;
 
 /** Shared behavior for producer and consumer. */
-final class SharedBehavior {
-  AdminClient client;
+final class SharedBehavior implements AutoCloseable {
+  private final AdminClient client;
 
   SharedBehavior(AdminClient client) {
     this.client = client;
@@ -57,4 +57,9 @@ List<PartitionInfo> partitionsFor(TopicPath topic, Duration timeout) {
       throw toKafka(t);
     }
   }
+
+  @Override
+  public void close() {
+    client.close();
+  }
 }
diff --git a/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumerTest.java b/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumerTest.java
@@ -30,7 +30,14 @@
 
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
-import com.google.cloud.pubsublite.*;
+import com.google.cloud.pubsublite.AdminClient;
+import com.google.cloud.pubsublite.CloudZone;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.ProjectNumber;
+import com.google.cloud.pubsublite.SubscriptionPath;
+import com.google.cloud.pubsublite.TopicName;
+import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.internal.CursorClient;
 import com.google.cloud.pubsublite.internal.testing.UnitTestExamples;
 import com.google.cloud.pubsublite.internal.wire.Assigner;

diff --git a/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteProducerTest.java b/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteProducerTest.java
@@ -26,7 +26,13 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.StatusCode.Code;
-import com.google.cloud.pubsublite.*;
+import com.google.cloud.pubsublite.AdminClient;
+import com.google.cloud.pubsublite.Message;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.PublishMetadata;
+import com.google.cloud.pubsublite.TopicName;
+import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.Publisher;
 import com.google.cloud.pubsublite.internal.testing.FakeApiService;

diff --git a/src/test/java/com/google/cloud/pubsublite/kafka/SharedBehaviorTest.java b/src/test/java/com/google/cloud/pubsublite/kafka/SharedBehaviorTest.java
@@ -19,6 +19,7 @@
 import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
@@ -73,4 +74,10 @@ public void partitionsForFailure() {
         BrokerNotAvailableException.class,
         () -> shared.partitionsFor(example(TopicPath.class), Duration.ofMillis(10)));
   }
+
+  @Test
+  public void close() {
+    shared.close();
+    verify(adminClient).close();
+  }
 }