Update tests to use DoFnRunner.ListOutputManager

----Release Notes----

[]
-------------
Created by MOE: http://code.google.com/p/moe-java
MOE_MIGRATED_REVID=99150521

Author: kennknowles

sdk/src/main/java/com/google/cloud/dataflow/sdk/transforms/DoFnTester.java

@@ -230,14 +230,15 @@ public void finishBundle() {
    */
   public List<OutputT> peekOutputElements() {
     // TODO: Should we return an unmodifiable list?
-    return Lists.transform(fnRunner.getReceiver(mainOutputTag),
-                           new Function<Object, OutputT>() {
-                             @Override
-                             @SuppressWarnings("unchecked")
-                             public OutputT apply(Object input) {
-                               return ((WindowedValue<OutputT>) input).getValue();
-                             }
-                           });
+    return Lists.transform(
+        outputManager.getOutput(mainOutputTag),
+        new Function<Object, OutputT>() {
+          @Override
+          @SuppressWarnings("unchecked")
+          public OutputT apply(Object input) {
+            return ((WindowedValue<OutputT>) input).getValue();
+          }
+        });
 
   }
 
@@ -272,12 +273,13 @@ public List<OutputT> takeOutputElements() {
    */
   public <T> List<T> peekSideOutputElements(TupleTag<T> tag) {
     // TODO: Should we return an unmodifiable list?
-    return Lists.transform(fnRunner.getReceiver(tag),
-                           new Function<Object, T>() {
-                             @Override
-                             public T apply(Object input) {
-                               return ((WindowedValue<T>) input).getValue();
-                             }});
+    return Lists.transform(
+        outputManager.getOutput(tag),
+        new Function<Object, T>() {
+          @Override
+          public T apply(Object input) {
+            return ((WindowedValue<T>) input).getValue();
+          }});
   }
 
   /**
@@ -323,6 +325,9 @@ enum State { UNSTARTED, STARTED, FINISHED }
   /** The original DoFn under test, if started. */
   DoFn<InputT, OutputT> fn;
 
+  /** The ListOutputManager to examine the outputs. */
+  DoFnRunner.ListOutputManager outputManager;
+
   /** The DoFnRunner if processing is in progress. */
   DoFnRunner<InputT, OutputT, List<WindowedValue<?>>> fnRunner;
 
@@ -341,6 +346,7 @@ enum State { UNSTARTED, STARTED, FINISHED }
 
   void resetState() {
     fn = null;
+    outputManager = null;
     fnRunner = null;
     counterSet = null;
     state = State.UNSTARTED;
@@ -358,11 +364,12 @@ void initializeState() {
         : sideInputs.entrySet()) {
       runnerSideInputs = runnerSideInputs.and(entry.getKey().getTagInternal(), entry.getValue());
     }
+    outputManager = new DoFnRunner.ListOutputManager();
     fnRunner = DoFnRunner.create(
         options,
         fn,
         DirectSideInputReader.of(runnerSideInputs),
-        new DoFnRunner.ListOutputManager(),
+        outputManager,
         mainOutputTag,
         sideOutputTags,
         DirectModeExecutionContext.create().createStepContext("stepName", "stepName"),

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/DoFnRunner.java

@@ -66,7 +66,6 @@ public interface OutputManager<ReceiverT> {
 
     /** Outputs a single element to the provided receiver. */
     public void output(ReceiverT receiver, WindowedValue<?> output);
-
   }
 
   /** The DoFn being run. */
@@ -197,11 +196,6 @@ public void finishBundle() {
     }
   }
 
-  /** Returns the receiver who gets outputs with the provided tag. */
-  public ReceiverT getReceiver(TupleTag<?> tag) {
-    return context.getReceiver(tag);
-  }
-
   /**
    * A concrete implementation of {@link DoFn.Context} used for running
    * a {@link DoFn}.

sdk/src/test/java/com/google/cloud/dataflow/sdk/util/GroupAlsoByWindowsDoFnTest.java

@@ -51,36 +51,39 @@
 
 /** Unit tests for {@link GroupAlsoByWindowsDoFn}. */
 @RunWith(JUnit4.class)
-@SuppressWarnings({"rawtypes", "unchecked"})
 public class GroupAlsoByWindowsDoFnTest {
   ExecutionContext execContext;
   CounterSet counters;
-  TupleTag<KV<String, Iterable<String>>> outputTag;
 
   @Before public void setUp() {
     execContext = new DirectModeExecutionContext();
     counters = new CounterSet();
-    outputTag = new TupleTag<>();
   }
 
   @Test public void testEmpty() throws Exception {
+    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
         KV<String, Iterable<String>>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
+        makeRunner(
+            outputTag, outputManager, WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
 
     runner.startBundle();
 
     runner.finishBundle();
 
-    List<KV<String, Iterable<String>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(0, result.size());
   }
 
   @Test public void testFixedWindows() throws Exception {
+    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
         KV<String, Iterable<String>>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
+        makeRunner(
+            outputTag, outputManager, WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
 
     runner.startBundle();
 
@@ -104,7 +107,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Iterable<String>>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(2, result.size());
 
@@ -124,10 +127,15 @@ public class GroupAlsoByWindowsDoFnTest {
   }
 
   @Test public void testSlidingWindows() throws Exception {
+    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
         KV<String, Iterable<String>>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(
-            SlidingWindows.of(Duration.millis(20)).every(Duration.millis(10))));
+        makeRunner(
+            outputTag,
+            outputManager,
+            WindowingStrategy.of(
+                SlidingWindows.of(Duration.millis(20)).every(Duration.millis(10))));
 
     runner.startBundle();
 
@@ -146,7 +154,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Iterable<String>>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(3, result.size());
 
@@ -173,10 +181,14 @@ public class GroupAlsoByWindowsDoFnTest {
   }
 
   @Test public void testSlidingWindowsCombine() throws Exception {
+    TupleTag<KV<String, Long>> outputTag = new TupleTag<>();
     CombineFn<Long, ?, Long> combineFn = new Sum.SumLongFn();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<Long>>>, KV<String, Long>,
         List<WindowedValue<?>>> runner =
         makeRunner(
+            outputTag,
+            outputManager,
             WindowingStrategy.of(SlidingWindows.of(Duration.millis(20)).every(Duration.millis(10))),
             combineFn.<String>asKeyedFn());
 
@@ -202,7 +214,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Long>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Long>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(3, result.size());
 
@@ -219,9 +231,12 @@ public class GroupAlsoByWindowsDoFnTest {
   }
 
   @Test public void testDiscontiguousWindows() throws Exception {
+    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
         KV<String, Iterable<String>>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
+        makeRunner(
+            outputTag, outputManager, WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
 
     runner.startBundle();
 
@@ -245,7 +260,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Iterable<String>>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(2, result.size());
 
@@ -265,9 +280,12 @@ public class GroupAlsoByWindowsDoFnTest {
   }
 
   @Test public void testSessions() throws Exception {
+    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
         KV<String, Iterable<String>>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10))));
+        makeRunner(outputTag, outputManager,
+            WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10))));
 
     runner.startBundle();
 
@@ -291,7 +309,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Iterable<String>>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
 
     assertEquals(2, result.size());
 
@@ -311,11 +329,16 @@ public class GroupAlsoByWindowsDoFnTest {
   }
 
   @Test public void testSessionsCombine() throws Exception {
+    TupleTag<KV<String, Long>> outputTag = new TupleTag<>();
     CombineFn<Long, ?, Long> combineFn = new Sum.SumLongFn();
+    DoFnRunner.ListOutputManager outputManager = new DoFnRunner.ListOutputManager();
     DoFnRunner<KV<String, Iterable<WindowedValue<Long>>>,
         KV<String, Long>, List<WindowedValue<?>>> runner =
-        makeRunner(WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10))),
-                   combineFn.<String>asKeyedFn());
+        makeRunner(
+            outputTag,
+            outputManager,
+            WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10))),
+            combineFn.<String>asKeyedFn());
     runner.startBundle();
 
     runner.processElement(WindowedValue.valueInEmptyWindows(
@@ -338,7 +361,7 @@ public class GroupAlsoByWindowsDoFnTest {
 
     runner.finishBundle();
 
-    List<WindowedValue<KV<String, Long>>> result = (List) runner.getReceiver(outputTag);
+    List<WindowedValue<KV<String, Long>>> result = outputManager.getOutput(outputTag);
 
     assertThat(result, Matchers.contains(
         WindowMatchers.isSingleWindowedValue(
@@ -349,40 +372,53 @@ public class GroupAlsoByWindowsDoFnTest {
             15, 15, 25)));
   }
 
-  private DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
-      KV<String, Iterable<String>>, List<WindowedValue<?>>> makeRunner(
+  private <ReceiverT>
+      DoFnRunner<KV<String, Iterable<WindowedValue<String>>>,
+          KV<String, Iterable<String>>, ReceiverT>
+      makeRunner(
+          TupleTag<KV<String, Iterable<String>>> outputTag,
+          DoFnRunner.OutputManager<ReceiverT> outputManager,
           WindowingStrategy<? super String, IntervalWindow> windowingStrategy) {
+
     GroupAlsoByWindowsDoFn<String, String, Iterable<String>, IntervalWindow> fn =
         GroupAlsoByWindowsDoFn.createForIterable(windowingStrategy, StringUtf8Coder.of());
-    return makeRunner(windowingStrategy, fn);
+
+    return makeRunner(outputTag, outputManager, windowingStrategy, fn);
   }
 
-  private DoFnRunner<KV<String, Iterable<WindowedValue<Long>>>,
-      KV<String, Long>, List<WindowedValue<?>>> makeRunner(
-        WindowingStrategy<? super String, IntervalWindow> windowingStrategy,
-        KeyedCombineFn<String, Long, ?, Long> combineFn) {
+  private <ReceiverT>
+      DoFnRunner<KV<String, Iterable<WindowedValue<Long>>>, KV<String, Long>, ReceiverT>
+      makeRunner(
+          TupleTag<KV<String, Long>> outputTag,
+          DoFnRunner.OutputManager<ReceiverT> outputManager,
+          WindowingStrategy<? super String, IntervalWindow> windowingStrategy,
+          KeyedCombineFn<String, Long, ?, Long> combineFn) {
+
     GroupAlsoByWindowsDoFn<String, Long, Long, IntervalWindow> fn =
-        GroupAlsoByWindowsDoFn.create(windowingStrategy, combineFn,
-            StringUtf8Coder.of(), BigEndianLongCoder.of());
+        GroupAlsoByWindowsDoFn.create(
+            windowingStrategy, combineFn, StringUtf8Coder.of(), BigEndianLongCoder.of());
 
-    return makeRunner(windowingStrategy, fn);
+    return makeRunner(outputTag, outputManager, windowingStrategy, fn);
   }
 
-  private <InputT, OutputT> DoFnRunner<KV<String, Iterable<WindowedValue<InputT>>>,
-    KV<String, OutputT>, List<WindowedValue<?>>> makeRunner(
+  private <InputT, OutputT, ReceiverT>
+      DoFnRunner<KV<String, Iterable<WindowedValue<InputT>>>, KV<String, OutputT>, ReceiverT>
+      makeRunner(
+        TupleTag<KV<String, OutputT>> outputTag,
+        DoFnRunner.OutputManager<ReceiverT> outputManager,
         WindowingStrategy<? super String, IntervalWindow> windowingStrategy,
         GroupAlsoByWindowsDoFn<String, InputT, OutputT, IntervalWindow> fn) {
-    return
-        DoFnRunner.create(
-            PipelineOptionsFactory.create(),
-            fn,
-            NullSideInputReader.empty(),
-            new DoFnRunner.ListOutputManager(),
-            (TupleTag<KV<String, OutputT>>) (TupleTag) outputTag,
-            new ArrayList<TupleTag<?>>(),
-            execContext.getStepContext("merge", "merge"),
-            counters.getAddCounterMutator(),
-            windowingStrategy);
+
+    return DoFnRunner.create(
+        PipelineOptionsFactory.create(),
+        fn,
+        NullSideInputReader.empty(),
+        outputManager,
+        outputTag,
+        new ArrayList<TupleTag<?>>(),
+        execContext.getStepContext("merge", "merge"),
+        counters.getAddCounterMutator(),
+        windowingStrategy);
   }
 
   private BoundedWindow window(long start, long end) {