fix(test): Fix flake in PubsubLiteSourceReaderTest (#44)

* fix: Fix flake in PubsubLiteSourceReaderTest * formatting
googleapis · Jul 29, 2021 · 09871f3 · 09871f3
1 parent 1c3a011
commit 09871f3
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/test/java/com/google/cloud/pubsublite/flink/reader/PubsubLiteSourceReaderTest.java b/src/test/java/com/google/cloud/pubsublite/flink/reader/PubsubLiteSourceReaderTest.java
@@ -109,6 +109,9 @@ public void testReader() throws Exception {
     assertThat(reader.pollNext(output)).isEqualTo(InputStatus.NOTHING_AVAILABLE);
 
     reader.notifyNoMoreSplits();
+    while (reader.pollNext(output) != InputStatus.END_OF_INPUT) {
+      assertThat(output.getEmittedRecords()).hasSize(6);
+    }
     assertThat(reader.pollNext(output)).isEqualTo(InputStatus.END_OF_INPUT);
     assertThat(output.getEmittedRecords()).containsExactly("0", "1", "2", "3", "5", "7");
   }