[FLINK-20625] Fix end-to-end test stopping condition

Author: RyanSkraba

flink-gcp-pubsub-e2e-tests/src/test/java/org/apache/flink/connector/gcp/pubsub/EmulatedPubSubSourceTest.java

@@ -25,7 +25,6 @@
 import org.apache.flink.api.common.serialization.SimpleStringSchema;
 import org.apache.flink.connector.gcp.pubsub.source.PubSubSource;
 import org.apache.flink.streaming.api.datastream.DataStream;
-import org.apache.flink.streaming.api.datastream.DataStreamUtils;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.connectors.gcp.pubsub.emulator.EmulatorCredentials;
 import org.apache.flink.streaming.connectors.gcp.pubsub.emulator.GCloudUnitTestBase;
@@ -126,16 +125,17 @@ public void testFlinkSource(boolean testWithFailure) throws Exception {
             fromPubSub = fromPubSub.map(new FailureMapFunction<>(3));
         }
 
-        List<String> output = new ArrayList<>();
-        DataStreamUtils.collect(fromPubSub).forEachRemaining(output::add);
+        // Asking for any more elements would wait forever, and there isn't a graceful way to
+        // indicate end of stream.
+        List<String> output = fromPubSub.executeAndCollect(input.size());
 
         assertEquals("Wrong number of elements", input.size(), output.size());
         for (String test : input) {
             assertTrue("Missing " + test, output.contains(test));
         }
     }
 
-    private class FailureMapFunction<T> extends RichMapFunction<T, T> {
+    private static class FailureMapFunction<T> extends RichMapFunction<T, T> {
         private final long numberOfRecordsUntilFailure;
         private long numberOfRecordsProcessed;