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
index cdad4024..2dd22a2f 100644
--- 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");
   }