[BEAM-22] Give root transforms step names

sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/ConsumerTrackingPipelineVisitor.java

@@ -76,12 +76,12 @@ public void leaveCompositeTransform(TransformTreeNode node) {
   public void visitTransform(TransformTreeNode node) {
     toFinalize.removeAll(node.getInput().expand());
     AppliedPTransform<?, ?, ?> appliedTransform = getAppliedTransform(node);
+    stepNames.put(appliedTransform, genStepName());
     if (node.getInput().expand().isEmpty()) {
       rootTransforms.add(appliedTransform);
     } else {
       for (PValue value : node.getInput().expand()) {
         valueToConsumers.get(value).add(appliedTransform);
-        stepNames.put(appliedTransform, genStepName());
       }
     }
   }

sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/runners/inprocess/ConsumerTrackingPipelineVisitorTest.java

@@ -176,6 +176,43 @@ public PDone apply(PInput input) {
     assertThat(visitor.getUnfinalizedPValues(), emptyIterable());
   }
 
+  @Test
+  public void getStepNamesContainsAllTransforms() {
+    PCollection<String> created = p.apply(Create.of("1", "2", "3"));
+    PCollection<String> transformed =
+        created.apply(
+            ParDo.of(
+                new DoFn<String, String>() {
+                  @Override
+                  public void processElement(DoFn<String, String>.ProcessContext c)
+                      throws Exception {
+                    c.output(Integer.toString(c.element().length()));
+                  }
+                }));
+    PDone finished =
+        transformed.apply(
+            new PTransform<PInput, PDone>() {
+              @Override
+              public PDone apply(PInput input) {
+                return PDone.in(input.getPipeline());
+              }
+            });
+
+    p.traverseTopologically(visitor);
+    assertThat(
+        visitor.getStepNames(),
+        Matchers.<AppliedPTransform<?, ?, ?>, String>hasEntry(
+            created.getProducingTransformInternal(), "s0"));
+    assertThat(
+        visitor.getStepNames(),
+        Matchers.<AppliedPTransform<?, ?, ?>, String>hasEntry(
+            transformed.getProducingTransformInternal(), "s1"));
+    assertThat(
+        visitor.getStepNames(),
+        Matchers.<AppliedPTransform<?, ?, ?>, String>hasEntry(
+            finished.getProducingTransformInternal(), "s2"));
+  }
+
   @Test
   public void traverseMultipleTimesThrows() {
     p.apply(Create.of(1, 2, 3));

sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/runners/inprocess/InProcessEvaluationContextTest.java

@@ -73,7 +73,6 @@
 
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -90,7 +89,8 @@ public class InProcessEvaluationContextTest {
   private PCollection<KV<String, Integer>> downstream;
   private PCollectionView<Iterable<Integer>> view;
   private PCollection<Long> unbounded;
-
+  private Collection<AppliedPTransform<?, ?, ?>> rootTransforms;
+  private Map<PValue, Collection<AppliedPTransform<?, ?, ?>>> valueToConsumers;
 
   @Before
   public void setup() {
@@ -103,31 +103,19 @@ public void setup() {
     downstream = created.apply(WithKeys.<String, Integer>of("foo"));
     view = created.apply(View.<Integer>asIterable());
     unbounded = p.apply(CountingInput.unbounded());
-    Collection<AppliedPTransform<?, ?, ?>> rootTransforms =
-        ImmutableList.<AppliedPTransform<?, ?, ?>>of(
-            created.getProducingTransformInternal(), unbounded.getProducingTransformInternal());
-    Map<PValue, Collection<AppliedPTransform<?, ?, ?>>> valueToConsumers = new HashMap<>();
-    valueToConsumers.put(
-        created,
-        ImmutableList.<AppliedPTransform<?, ?, ?>>of(
-            downstream.getProducingTransformInternal(), view.getProducingTransformInternal()));
-    valueToConsumers.put(unbounded, ImmutableList.<AppliedPTransform<?, ?, ?>>of());
-    valueToConsumers.put(downstream, ImmutableList.<AppliedPTransform<?, ?, ?>>of());
-    valueToConsumers.put(view, ImmutableList.<AppliedPTransform<?, ?, ?>>of());
-
-    Map<AppliedPTransform<?, ?, ?>, String> stepNames = new HashMap<>();
-    stepNames.put(created.getProducingTransformInternal(), "s1");
-    stepNames.put(downstream.getProducingTransformInternal(), "s2");
-    stepNames.put(view.getProducingTransformInternal(), "s3");
-    stepNames.put(unbounded.getProducingTransformInternal(), "s4");
-
-    Collection<PCollectionView<?>> views = ImmutableList.<PCollectionView<?>>of(view);
-    context = InProcessEvaluationContext.create(
+
+    ConsumerTrackingPipelineVisitor cVis = new ConsumerTrackingPipelineVisitor();
+    p.traverseTopologically(cVis);
+    rootTransforms = cVis.getRootTransforms();
+    valueToConsumers = cVis.getValueToConsumers();
+
+    context =
+        InProcessEvaluationContext.create(
             runner.getPipelineOptions(),
             rootTransforms,
             valueToConsumers,
-            stepNames,
-            views);
+            cVis.getStepNames(),
+            cVis.getViews());
   }
 
   @Test
@@ -492,16 +480,14 @@ public void isDoneWithPartiallyDone() {
         null,
         ImmutableList.<TimerData>of(),
         StepTransformResult.withoutHold(unbounded.getProducingTransformInternal()).build());
-    context.handleResult(
-        committedBundle,
-        ImmutableList.<TimerData>of(),
-        StepTransformResult.withoutHold(downstream.getProducingTransformInternal()).build());
     assertThat(context.isDone(), is(false));
 
-    context.handleResult(
-        committedBundle,
-        ImmutableList.<TimerData>of(),
-        StepTransformResult.withoutHold(view.getProducingTransformInternal()).build());
+    for (AppliedPTransform<?, ?, ?> consumers : valueToConsumers.get(created)) {
+      context.handleResult(
+          committedBundle,
+          ImmutableList.<TimerData>of(),
+          StepTransformResult.withoutHold(consumers).build());
+    }
     assertThat(context.isDone(), is(true));
   }