fix: Fail hard in subscribe() call on assignment failure (#197)

googleapis · Aug 16, 2021 · 4d15dcc · 4d15dcc
1 parent 33f452c
commit 4d15dcc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/main/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumer.java b/src/main/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumer.java
@@ -209,7 +209,7 @@ public void subscribe(
     consumer = Optional.of(consumerFactory.newConsumer());
     try {
       assigner = Optional.of(assignerFactory.New(newAssignmentReceiver(consumerRebalanceListener)));
-      assigner.get().startAsync();
+      assigner.get().startAsync().awaitRunning();
     } catch (ApiException e) {
       throw toKafka(e);
     }

diff --git a/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumerTest.java b/src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumerTest.java
@@ -42,6 +42,7 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.CursorClient;
 import com.google.cloud.pubsublite.internal.TopicStatsClient;
+import com.google.cloud.pubsublite.internal.testing.FakeApiService;
 import com.google.cloud.pubsublite.internal.testing.UnitTestExamples;
 import com.google.cloud.pubsublite.internal.wire.Assigner;
 import com.google.cloud.pubsublite.internal.wire.AssignerFactory;
@@ -80,6 +81,7 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mock;
+import org.mockito.Spy;
 
 @RunWith(JUnit4.class)
 public class PubsubLiteConsumerTest {
@@ -112,7 +114,9 @@ private static <T> T example(Class<T> klass) {
   @Mock AdminClient adminClient;
   @Mock TopicStatsClient topicStatsClient;
 
-  @Mock Assigner assigner;
+  abstract static class FakeAssigner extends FakeApiService implements Assigner {}
+
+  @Spy FakeAssigner assigner;
   @Mock SingleSubscriptionConsumer underlying;
 
   Consumer<byte[], byte[]> consumer;
@@ -305,6 +309,7 @@ public void validSubscribe() throws Exception {
     verify(consumerFactory).newConsumer();
     verify(assignerFactory).New(any());
     verify(assigner).startAsync();
+    verify(assigner).awaitRunning();
     receiver.get().handleAssignment(ImmutableSet.of(Partition.of(5)));
     verify(listener)
         .onPartitionsAssigned(