From 09871f3797350f80023aaba05560adb6924ae75c Mon Sep 17 00:00:00 2001
From: palmere-google <68394592+palmere-google@users.noreply.github.com>
Date: Thu, 29 Jul 2021 13:33:24 -0400
Subject: [PATCH] fix(test): Fix flake in PubsubLiteSourceReaderTest (#44)

* fix: Fix flake in PubsubLiteSourceReaderTest

* formatting
---
 .../pubsublite/flink/reader/PubsubLiteSourceReaderTest.java    | 3 +++
 1 file changed, 3 insertions(+)

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");
   }