You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ke...@apache.org on 2016/06/16 02:30:42 UTC
[07/12] incubator-beam git commit: Remove InProcess Prefixes
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectTimerInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectTimerInternalsTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectTimerInternalsTest.java
new file mode 100644
index 0000000..c038910
--- /dev/null
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectTimerInternalsTest.java
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.direct;
+
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.when;
+
+import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
+import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate.TimerUpdateBuilder;
+import org.apache.beam.runners.direct.WatermarkManager.TransformWatermarks;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.util.TimeDomain;
+import org.apache.beam.sdk.util.TimerInternals.TimerData;
+import org.apache.beam.sdk.util.state.StateNamespaces;
+
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link DirectTimerInternals}.
+ */
+@RunWith(JUnit4.class)
+public class DirectTimerInternalsTest {
+ private MockClock clock;
+ @Mock private TransformWatermarks watermarks;
+
+ private TimerUpdateBuilder timerUpdateBuilder;
+
+ private DirectTimerInternals internals;
+
+ @Before
+ public void setup() {
+ MockitoAnnotations.initMocks(this);
+ clock = MockClock.fromInstant(new Instant(0));
+
+ timerUpdateBuilder = TimerUpdate.builder(StructuralKey.of(1234, VarIntCoder.of()));
+
+ internals = DirectTimerInternals.create(clock, watermarks, timerUpdateBuilder);
+ }
+
+ @Test
+ public void setTimerAddsToBuilder() {
+ TimerData eventTimer =
+ TimerData.of(StateNamespaces.global(), new Instant(20145L), TimeDomain.EVENT_TIME);
+ TimerData processingTimer =
+ TimerData.of(StateNamespaces.global(), new Instant(125555555L), TimeDomain.PROCESSING_TIME);
+ TimerData synchronizedProcessingTimer =
+ TimerData.of(
+ StateNamespaces.global(),
+ new Instant(98745632189L),
+ TimeDomain.SYNCHRONIZED_PROCESSING_TIME);
+ internals.setTimer(eventTimer);
+ internals.setTimer(processingTimer);
+ internals.setTimer(synchronizedProcessingTimer);
+
+ assertThat(
+ internals.getTimerUpdate().getSetTimers(),
+ containsInAnyOrder(eventTimer, synchronizedProcessingTimer, processingTimer));
+ }
+
+ @Test
+ public void deleteTimerDeletesOnBuilder() {
+ TimerData eventTimer =
+ TimerData.of(StateNamespaces.global(), new Instant(20145L), TimeDomain.EVENT_TIME);
+ TimerData processingTimer =
+ TimerData.of(StateNamespaces.global(), new Instant(125555555L), TimeDomain.PROCESSING_TIME);
+ TimerData synchronizedProcessingTimer =
+ TimerData.of(
+ StateNamespaces.global(),
+ new Instant(98745632189L),
+ TimeDomain.SYNCHRONIZED_PROCESSING_TIME);
+ internals.deleteTimer(eventTimer);
+ internals.deleteTimer(processingTimer);
+ internals.deleteTimer(synchronizedProcessingTimer);
+
+ assertThat(
+ internals.getTimerUpdate().getDeletedTimers(),
+ containsInAnyOrder(eventTimer, synchronizedProcessingTimer, processingTimer));
+ }
+
+ @Test
+ public void getProcessingTimeIsClockNow() {
+ assertThat(internals.currentProcessingTime(), equalTo(clock.now()));
+ Instant oldProcessingTime = internals.currentProcessingTime();
+
+ clock.advance(Duration.standardHours(12));
+
+ assertThat(internals.currentProcessingTime(), equalTo(clock.now()));
+ assertThat(
+ internals.currentProcessingTime(),
+ equalTo(oldProcessingTime.plus(Duration.standardHours(12))));
+ }
+
+ @Test
+ public void getSynchronizedProcessingTimeIsWatermarkSynchronizedInputTime() {
+ when(watermarks.getSynchronizedProcessingInputTime()).thenReturn(new Instant(12345L));
+ assertThat(internals.currentSynchronizedProcessingTime(), equalTo(new Instant(12345L)));
+ }
+
+ @Test
+ public void getInputWatermarkTimeUsesWatermarkTime() {
+ when(watermarks.getInputWatermark()).thenReturn(new Instant(8765L));
+ assertThat(internals.currentInputWatermarkTime(), equalTo(new Instant(8765L)));
+ }
+
+ @Test
+ public void getOutputWatermarkTimeUsesWatermarkTime() {
+ when(watermarks.getOutputWatermark()).thenReturn(new Instant(25525L));
+ assertThat(internals.currentOutputWatermarkTime(), equalTo(new Instant(25525L)));
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EncodabilityEnforcementFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EncodabilityEnforcementFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EncodabilityEnforcementFactoryTest.java
index e129489..b903ef1 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EncodabilityEnforcementFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EncodabilityEnforcementFactoryTest.java
@@ -51,7 +51,7 @@ import java.util.Collections;
public class EncodabilityEnforcementFactoryTest {
@Rule public ExpectedException thrown = ExpectedException.none();
private EncodabilityEnforcementFactory factory = EncodabilityEnforcementFactory.create();
- private BundleFactory bundleFactory = InProcessBundleFactory.create();
+ private BundleFactory bundleFactory = ImmutableListBundleFactory.create();
@Test
public void encodeFailsThrows() {
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
new file mode 100644
index 0000000..1726866
--- /dev/null
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
@@ -0,0 +1,545 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.direct;
+
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.emptyIterable;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
+import static org.junit.Assert.assertThat;
+
+import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext;
+import org.apache.beam.runners.direct.DirectRunner.CommittedBundle;
+import org.apache.beam.runners.direct.DirectRunner.PCollectionViewWriter;
+import org.apache.beam.runners.direct.DirectRunner.UncommittedBundle;
+import org.apache.beam.runners.direct.WatermarkManager.FiredTimers;
+import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.io.CountingInput;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.AppliedPTransform;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
+import org.apache.beam.sdk.util.SideInputReader;
+import org.apache.beam.sdk.util.TimeDomain;
+import org.apache.beam.sdk.util.TimerInternals.TimerData;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.WindowingStrategy;
+import org.apache.beam.sdk.util.common.Counter;
+import org.apache.beam.sdk.util.common.Counter.AggregationKind;
+import org.apache.beam.sdk.util.common.CounterSet;
+import org.apache.beam.sdk.util.state.BagState;
+import org.apache.beam.sdk.util.state.CopyOnAccessInMemoryStateInternals;
+import org.apache.beam.sdk.util.state.StateNamespaces;
+import org.apache.beam.sdk.util.state.StateTag;
+import org.apache.beam.sdk.util.state.StateTags;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollection.IsBounded;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PValue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+
+import org.hamcrest.Matchers;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Tests for {@link EvaluationContext}.
+ */
+@RunWith(JUnit4.class)
+public class EvaluationContextTest {
+ private TestPipeline p;
+ private EvaluationContext context;
+
+ private PCollection<Integer> created;
+ 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;
+
+ private BundleFactory bundleFactory;
+
+ @Before
+ public void setup() {
+ DirectRunner runner =
+ DirectRunner.fromOptions(PipelineOptionsFactory.create());
+
+ p = TestPipeline.create();
+
+ created = p.apply(Create.of(1, 2, 3));
+ downstream = created.apply(WithKeys.<String, Integer>of("foo"));
+ view = created.apply(View.<Integer>asIterable());
+ unbounded = p.apply(CountingInput.unbounded());
+
+ ConsumerTrackingPipelineVisitor cVis = new ConsumerTrackingPipelineVisitor();
+ p.traverseTopologically(cVis);
+ rootTransforms = cVis.getRootTransforms();
+ valueToConsumers = cVis.getValueToConsumers();
+
+ bundleFactory = ImmutableListBundleFactory.create();
+
+ context =
+ EvaluationContext.create(
+ runner.getPipelineOptions(),
+ ImmutableListBundleFactory.create(),
+ rootTransforms,
+ valueToConsumers,
+ cVis.getStepNames(),
+ cVis.getViews());
+ }
+
+ @Test
+ public void writeToViewWriterThenReadReads() {
+ PCollectionViewWriter<Integer, Iterable<Integer>> viewWriter =
+ context.createPCollectionViewWriter(
+ PCollection.<Iterable<Integer>>createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED),
+ view);
+ BoundedWindow window = new TestBoundedWindow(new Instant(1024L));
+ BoundedWindow second = new TestBoundedWindow(new Instant(899999L));
+ WindowedValue<Integer> firstValue =
+ WindowedValue.of(1, new Instant(1222), window, PaneInfo.ON_TIME_AND_ONLY_FIRING);
+ WindowedValue<Integer> secondValue =
+ WindowedValue.of(
+ 2, new Instant(8766L), second, PaneInfo.createPane(true, false, Timing.ON_TIME, 0, 0));
+ Iterable<WindowedValue<Integer>> values = ImmutableList.of(firstValue, secondValue);
+ viewWriter.add(values);
+
+ SideInputReader reader =
+ context.createSideInputReader(ImmutableList.<PCollectionView<?>>of(view));
+ assertThat(reader.get(view, window), containsInAnyOrder(1));
+ assertThat(reader.get(view, second), containsInAnyOrder(2));
+
+ WindowedValue<Integer> overrittenSecondValue =
+ WindowedValue.of(
+ 4444, new Instant(8677L), second, PaneInfo.createPane(false, true, Timing.LATE, 1, 1));
+ viewWriter.add(Collections.singleton(overrittenSecondValue));
+ assertThat(reader.get(view, second), containsInAnyOrder(2));
+ // The cached value is served in the earlier reader
+ reader = context.createSideInputReader(ImmutableList.<PCollectionView<?>>of(view));
+ assertThat(reader.get(view, second), containsInAnyOrder(4444));
+ }
+
+ @Test
+ public void getExecutionContextSameStepSameKeyState() {
+ DirectExecutionContext fooContext =
+ context.getExecutionContext(created.getProducingTransformInternal(),
+ StructuralKey.of("foo", StringUtf8Coder.of()));
+
+ StateTag<Object, BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
+
+ DirectStepContext stepContext = fooContext.getOrCreateStepContext("s1", "s1");
+ stepContext.stateInternals().state(StateNamespaces.global(), intBag).add(1);
+
+ context.handleResult(ImmutableListBundleFactory.create()
+ .createKeyedBundle(null, StructuralKey.of("foo", StringUtf8Coder.of()), created)
+ .commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(created.getProducingTransformInternal())
+ .withState(stepContext.commitState())
+ .build());
+
+ DirectExecutionContext secondFooContext =
+ context.getExecutionContext(created.getProducingTransformInternal(),
+ StructuralKey.of("foo", StringUtf8Coder.of()));
+ assertThat(
+ secondFooContext
+ .getOrCreateStepContext("s1", "s1")
+ .stateInternals()
+ .state(StateNamespaces.global(), intBag)
+ .read(),
+ contains(1));
+ }
+
+
+ @Test
+ public void getExecutionContextDifferentKeysIndependentState() {
+ DirectExecutionContext fooContext =
+ context.getExecutionContext(created.getProducingTransformInternal(),
+ StructuralKey.of("foo", StringUtf8Coder.of()));
+
+ StateTag<Object, BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
+
+ fooContext
+ .getOrCreateStepContext("s1", "s1")
+ .stateInternals()
+ .state(StateNamespaces.global(), intBag)
+ .add(1);
+
+ DirectExecutionContext barContext =
+ context.getExecutionContext(created.getProducingTransformInternal(),
+ StructuralKey.of("bar", StringUtf8Coder.of()));
+ assertThat(barContext, not(equalTo(fooContext)));
+ assertThat(
+ barContext
+ .getOrCreateStepContext("s1", "s1")
+ .stateInternals()
+ .state(StateNamespaces.global(), intBag)
+ .read(),
+ emptyIterable());
+ }
+
+ @Test
+ public void getExecutionContextDifferentStepsIndependentState() {
+ StructuralKey<?> myKey = StructuralKey.of("foo", StringUtf8Coder.of());
+ DirectExecutionContext fooContext =
+ context.getExecutionContext(created.getProducingTransformInternal(), myKey);
+
+ StateTag<Object, BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
+
+ fooContext
+ .getOrCreateStepContext("s1", "s1")
+ .stateInternals()
+ .state(StateNamespaces.global(), intBag)
+ .add(1);
+
+ DirectExecutionContext barContext =
+ context.getExecutionContext(downstream.getProducingTransformInternal(), myKey);
+ assertThat(
+ barContext
+ .getOrCreateStepContext("s1", "s1")
+ .stateInternals()
+ .state(StateNamespaces.global(), intBag)
+ .read(),
+ emptyIterable());
+ }
+
+ @Test
+ public void handleResultMergesCounters() {
+ CounterSet counters = context.createCounterSet();
+ Counter<Long> myCounter = Counter.longs("foo", AggregationKind.SUM);
+ counters.addCounter(myCounter);
+
+ myCounter.addValue(4L);
+ TransformResult result =
+ StepTransformResult.withoutHold(created.getProducingTransformInternal())
+ .withCounters(counters)
+ .build();
+ context.handleResult(null, ImmutableList.<TimerData>of(), result);
+ assertThat((Long) context.getCounters().getExistingCounter("foo").getAggregate(), equalTo(4L));
+
+ CounterSet againCounters = context.createCounterSet();
+ Counter<Long> myLongCounterAgain = Counter.longs("foo", AggregationKind.SUM);
+ againCounters.add(myLongCounterAgain);
+ myLongCounterAgain.addValue(8L);
+
+ TransformResult secondResult =
+ StepTransformResult.withoutHold(downstream.getProducingTransformInternal())
+ .withCounters(againCounters)
+ .build();
+ context.handleResult(
+ context.createRootBundle(created).commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ secondResult);
+ assertThat((Long) context.getCounters().getExistingCounter("foo").getAggregate(), equalTo(12L));
+ }
+
+ @Test
+ public void handleResultStoresState() {
+ StructuralKey<?> myKey = StructuralKey.of("foo".getBytes(), ByteArrayCoder.of());
+ DirectExecutionContext fooContext =
+ context.getExecutionContext(downstream.getProducingTransformInternal(), myKey);
+
+ StateTag<Object, BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
+
+ CopyOnAccessInMemoryStateInternals<Object> state =
+ fooContext.getOrCreateStepContext("s1", "s1").stateInternals();
+ BagState<Integer> bag = state.state(StateNamespaces.global(), intBag);
+ bag.add(1);
+ bag.add(2);
+ bag.add(4);
+
+ TransformResult stateResult =
+ StepTransformResult.withoutHold(downstream.getProducingTransformInternal())
+ .withState(state)
+ .build();
+
+ context.handleResult(
+ context.createKeyedBundle(null, myKey, created).commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ stateResult);
+
+ DirectExecutionContext afterResultContext =
+ context.getExecutionContext(downstream.getProducingTransformInternal(), myKey);
+
+ CopyOnAccessInMemoryStateInternals<Object> afterResultState =
+ afterResultContext.getOrCreateStepContext("s1", "s1").stateInternals();
+ assertThat(afterResultState.state(StateNamespaces.global(), intBag).read(), contains(1, 2, 4));
+ }
+
+ @Test
+ public void callAfterOutputMustHaveBeenProducedAfterEndOfWatermarkCallsback() throws Exception {
+ final CountDownLatch callLatch = new CountDownLatch(1);
+ Runnable callback =
+ new Runnable() {
+ @Override
+ public void run() {
+ callLatch.countDown();
+ }
+ };
+
+ // Should call back after the end of the global window
+ context.scheduleAfterOutputWouldBeProduced(
+ downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);
+
+ TransformResult result =
+ StepTransformResult.withHold(created.getProducingTransformInternal(), new Instant(0))
+ .build();
+
+ context.handleResult(null, ImmutableList.<TimerData>of(), result);
+ // Difficult to demonstrate that we took no action in a multithreaded world; poll for a bit
+ // will likely be flaky if this logic is broken
+ assertThat(callLatch.await(500L, TimeUnit.MILLISECONDS), is(false));
+
+ TransformResult finishedResult =
+ StepTransformResult.withoutHold(created.getProducingTransformInternal()).build();
+ context.handleResult(null, ImmutableList.<TimerData>of(), finishedResult);
+ context.forceRefresh();
+ // Obtain the value via blocking call
+ assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
+ }
+
+ @Test
+ public void callAfterOutputMustHaveBeenProducedAlreadyAfterCallsImmediately() throws Exception {
+ TransformResult finishedResult =
+ StepTransformResult.withoutHold(created.getProducingTransformInternal()).build();
+ context.handleResult(null, ImmutableList.<TimerData>of(), finishedResult);
+
+ final CountDownLatch callLatch = new CountDownLatch(1);
+ context.extractFiredTimers();
+ Runnable callback =
+ new Runnable() {
+ @Override
+ public void run() {
+ callLatch.countDown();
+ }
+ };
+ context.scheduleAfterOutputWouldBeProduced(
+ downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);
+ assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
+ }
+
+ @Test
+ public void extractFiredTimersExtractsTimers() {
+ TransformResult holdResult =
+ StepTransformResult.withHold(created.getProducingTransformInternal(), new Instant(0))
+ .build();
+ context.handleResult(null, ImmutableList.<TimerData>of(), holdResult);
+
+ StructuralKey<?> key = StructuralKey.of("foo".length(), VarIntCoder.of());
+ TimerData toFire =
+ TimerData.of(StateNamespaces.global(), new Instant(100L), TimeDomain.EVENT_TIME);
+ TransformResult timerResult =
+ StepTransformResult.withoutHold(downstream.getProducingTransformInternal())
+ .withState(CopyOnAccessInMemoryStateInternals.withUnderlying(key, null))
+ .withTimerUpdate(TimerUpdate.builder(key).setTimer(toFire).build())
+ .build();
+
+ // haven't added any timers, must be empty
+ assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+ context.handleResult(
+ context.createKeyedBundle(null, key, created).commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ timerResult);
+
+ // timer hasn't fired
+ assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+
+ TransformResult advanceResult =
+ StepTransformResult.withoutHold(created.getProducingTransformInternal()).build();
+ // Should cause the downstream timer to fire
+ context.handleResult(null, ImmutableList.<TimerData>of(), advanceResult);
+
+ Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>> fired =
+ context.extractFiredTimers();
+ assertThat(
+ fired,
+ Matchers.<AppliedPTransform<?, ?, ?>>hasKey(downstream.getProducingTransformInternal()));
+ Map<StructuralKey<?>, FiredTimers> downstreamFired =
+ fired.get(downstream.getProducingTransformInternal());
+ assertThat(downstreamFired, Matchers.<Object>hasKey(key));
+
+ FiredTimers firedForKey = downstreamFired.get(key);
+ assertThat(firedForKey.getTimers(TimeDomain.PROCESSING_TIME), emptyIterable());
+ assertThat(firedForKey.getTimers(TimeDomain.SYNCHRONIZED_PROCESSING_TIME), emptyIterable());
+ assertThat(firedForKey.getTimers(TimeDomain.EVENT_TIME), contains(toFire));
+
+ // Don't reextract timers
+ assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+ }
+
+ @Test
+ public void createBundleKeyedResultPropagatesKey() {
+ StructuralKey<String> key = StructuralKey.of("foo", StringUtf8Coder.of());
+ CommittedBundle<KV<String, Integer>> newBundle =
+ context
+ .createBundle(
+ bundleFactory.createKeyedBundle(
+ null, key,
+ created).commit(Instant.now()),
+ downstream).commit(Instant.now());
+ assertThat(newBundle.getKey(), Matchers.<StructuralKey<?>>equalTo(key));
+ }
+
+ @Test
+ public void createKeyedBundleKeyed() {
+ StructuralKey<String> key = StructuralKey.of("foo", StringUtf8Coder.of());
+ CommittedBundle<KV<String, Integer>> keyedBundle =
+ context.createKeyedBundle(
+ bundleFactory.createRootBundle(created).commit(Instant.now()),
+ key,
+ downstream).commit(Instant.now());
+ assertThat(keyedBundle.getKey(),
+ Matchers.<StructuralKey<?>>equalTo(key));
+ }
+
+ @Test
+ public void isDoneWithUnboundedPCollectionAndShutdown() {
+ context.getPipelineOptions().setShutdownUnboundedProducersWithMaxWatermark(true);
+ assertThat(context.isDone(unbounded.getProducingTransformInternal()), is(false));
+
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(unbounded.getProducingTransformInternal()).build());
+ context.extractFiredTimers();
+ assertThat(context.isDone(unbounded.getProducingTransformInternal()), is(true));
+ }
+
+ @Test
+ public void isDoneWithUnboundedPCollectionAndNotShutdown() {
+ context.getPipelineOptions().setShutdownUnboundedProducersWithMaxWatermark(false);
+ assertThat(context.isDone(unbounded.getProducingTransformInternal()), is(false));
+
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(unbounded.getProducingTransformInternal()).build());
+ assertThat(context.isDone(unbounded.getProducingTransformInternal()), is(false));
+ }
+
+ @Test
+ public void isDoneWithOnlyBoundedPCollections() {
+ context.getPipelineOptions().setShutdownUnboundedProducersWithMaxWatermark(false);
+ assertThat(context.isDone(created.getProducingTransformInternal()), is(false));
+
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(created.getProducingTransformInternal()).build());
+ context.extractFiredTimers();
+ assertThat(context.isDone(created.getProducingTransformInternal()), is(true));
+ }
+
+ @Test
+ public void isDoneWithPartiallyDone() {
+ context.getPipelineOptions().setShutdownUnboundedProducersWithMaxWatermark(true);
+ assertThat(context.isDone(), is(false));
+
+ UncommittedBundle<Integer> rootBundle = context.createRootBundle(created);
+ rootBundle.add(WindowedValue.valueInGlobalWindow(1));
+ CommittedResult handleResult =
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(created.getProducingTransformInternal())
+ .addOutput(rootBundle)
+ .build());
+ @SuppressWarnings("unchecked")
+ CommittedBundle<Integer> committedBundle =
+ (CommittedBundle<Integer>) Iterables.getOnlyElement(handleResult.getOutputs());
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(unbounded.getProducingTransformInternal()).build());
+ assertThat(context.isDone(), is(false));
+
+ for (AppliedPTransform<?, ?, ?> consumers : valueToConsumers.get(created)) {
+ context.handleResult(
+ committedBundle,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(consumers).build());
+ }
+ context.extractFiredTimers();
+ assertThat(context.isDone(), is(true));
+ }
+
+ @Test
+ public void isDoneWithUnboundedAndNotShutdown() {
+ context.getPipelineOptions().setShutdownUnboundedProducersWithMaxWatermark(false);
+ assertThat(context.isDone(), is(false));
+
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(created.getProducingTransformInternal()).build());
+ context.handleResult(
+ null,
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(unbounded.getProducingTransformInternal()).build());
+ context.handleResult(
+ context.createRootBundle(created).commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(downstream.getProducingTransformInternal()).build());
+ context.extractFiredTimers();
+ assertThat(context.isDone(), is(false));
+
+ context.handleResult(
+ context.createRootBundle(created).commit(Instant.now()),
+ ImmutableList.<TimerData>of(),
+ StepTransformResult.withoutHold(view.getProducingTransformInternal()).build());
+ context.extractFiredTimers();
+ assertThat(context.isDone(), is(false));
+ }
+
+ private static class TestBoundedWindow extends BoundedWindow {
+ private final Instant ts;
+
+ public TestBoundedWindow(Instant ts) {
+ this.ts = ts;
+ }
+
+ @Override
+ public Instant maxTimestamp() {
+ return ts;
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/FlattenEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/FlattenEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/FlattenEvaluatorFactoryTest.java
index 5efb090..0bc3036 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/FlattenEvaluatorFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/FlattenEvaluatorFactoryTest.java
@@ -45,7 +45,7 @@ import org.junit.runners.JUnit4;
*/
@RunWith(JUnit4.class)
public class FlattenEvaluatorFactoryTest {
- private BundleFactory bundleFactory = InProcessBundleFactory.create();
+ private BundleFactory bundleFactory = ImmutableListBundleFactory.create();
@Test
public void testFlattenInMemoryEvaluator() throws Exception {
TestPipeline p = TestPipeline.create();
@@ -60,7 +60,7 @@ public class FlattenEvaluatorFactoryTest {
CommittedBundle<Integer> rightBundle =
bundleFactory.createRootBundle(right).commit(Instant.now());
- InProcessEvaluationContext context = mock(InProcessEvaluationContext.class);
+ EvaluationContext context = mock(EvaluationContext.class);
UncommittedBundle<Integer> flattenedLeftBundle = bundleFactory.createRootBundle(flattened);
UncommittedBundle<Integer> flattenedRightBundle = bundleFactory.createRootBundle(flattened);
@@ -87,8 +87,8 @@ public class FlattenEvaluatorFactoryTest {
rightSideEvaluator.processElement(
WindowedValue.timestampedValueInGlobalWindow(-4, new Instant(-4096)));
- InProcessTransformResult rightSideResult = rightSideEvaluator.finishBundle();
- InProcessTransformResult leftSideResult = leftSideEvaluator.finishBundle();
+ TransformResult rightSideResult = rightSideEvaluator.finishBundle();
+ TransformResult leftSideResult = leftSideEvaluator.finishBundle();
assertThat(
rightSideResult.getOutputBundles(),
@@ -124,13 +124,13 @@ public class FlattenEvaluatorFactoryTest {
PCollection<Integer> flattened = list.apply(Flatten.<Integer>pCollections());
- InProcessEvaluationContext context = mock(InProcessEvaluationContext.class);
+ EvaluationContext context = mock(EvaluationContext.class);
FlattenEvaluatorFactory factory = new FlattenEvaluatorFactory();
TransformEvaluator<Integer> emptyEvaluator =
factory.forApplication(flattened.getProducingTransformInternal(), null, context);
- InProcessTransformResult leftSideResult = emptyEvaluator.finishBundle();
+ TransformResult leftSideResult = emptyEvaluator.finishBundle();
assertThat(leftSideResult.getOutputBundles(), emptyIterable());
assertThat(
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyEvaluatorFactoryTest.java
index b589db0..a529043 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyEvaluatorFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyEvaluatorFactoryTest.java
@@ -22,6 +22,7 @@ import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupByKeyOnly;
import org.apache.beam.runners.direct.DirectRunner.CommittedBundle;
import org.apache.beam.runners.direct.DirectRunner.UncommittedBundle;
import org.apache.beam.sdk.coders.Coder;
@@ -52,7 +53,7 @@ import org.junit.runners.JUnit4;
*/
@RunWith(JUnit4.class)
public class GroupByKeyEvaluatorFactoryTest {
- private BundleFactory bundleFactory = InProcessBundleFactory.create();
+ private BundleFactory bundleFactory = ImmutableListBundleFactory.create();
@Test
public void testInMemoryEvaluator() throws Exception {
@@ -68,11 +69,11 @@ public class GroupByKeyEvaluatorFactoryTest {
PCollection<KV<String, WindowedValue<Integer>>> kvs =
values.apply(new ReifyTimestampsAndWindows<String, Integer>());
PCollection<KeyedWorkItem<String, Integer>> groupedKvs =
- kvs.apply(new InProcessGroupByKey.InProcessGroupByKeyOnly<String, Integer>());
+ kvs.apply(new DirectGroupByKeyOnly<String, Integer>());
CommittedBundle<KV<String, WindowedValue<Integer>>> inputBundle =
bundleFactory.createRootBundle(kvs).commit(Instant.now());
- InProcessEvaluationContext evaluationContext = mock(InProcessEvaluationContext.class);
+ EvaluationContext evaluationContext = mock(EvaluationContext.class);
StructuralKey<String> fooKey = StructuralKey.of("foo", StringUtf8Coder.of());
UncommittedBundle<KeyedWorkItem<String, Integer>> fooBundle =
bundleFactory.createKeyedBundle(null, fooKey, groupedKvs);
@@ -100,7 +101,7 @@ public class GroupByKeyEvaluatorFactoryTest {
Coder<String> keyCoder =
((KvCoder<String, WindowedValue<Integer>>) kvs.getCoder()).getKeyCoder();
TransformEvaluator<KV<String, WindowedValue<Integer>>> evaluator =
- new InProcessGroupByKeyOnlyEvaluatorFactory()
+ new GroupByKeyOnlyEvaluatorFactory()
.forApplication(
groupedKvs.getProducingTransformInternal(), inputBundle, evaluationContext);
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactoryTest.java
new file mode 100644
index 0000000..5fb5a76
--- /dev/null
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactoryTest.java
@@ -0,0 +1,197 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.direct;
+
+import static org.hamcrest.Matchers.contains;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupByKeyOnly;
+import org.apache.beam.runners.direct.DirectRunner.CommittedBundle;
+import org.apache.beam.runners.direct.DirectRunner.UncommittedBundle;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.util.GroupByKeyViaGroupByKeyOnly.ReifyTimestampsAndWindows;
+import org.apache.beam.sdk.util.KeyedWorkItem;
+import org.apache.beam.sdk.util.KeyedWorkItems;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+
+import com.google.common.collect.HashMultiset;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multiset;
+
+import org.hamcrest.BaseMatcher;
+import org.hamcrest.Description;
+import org.joda.time.Instant;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests for {@link GroupByKeyOnlyEvaluatorFactory}.
+ */
+@RunWith(JUnit4.class)
+public class GroupByKeyOnlyEvaluatorFactoryTest {
+ private BundleFactory bundleFactory = ImmutableListBundleFactory.create();
+
+ @Test
+ public void testInMemoryEvaluator() throws Exception {
+ TestPipeline p = TestPipeline.create();
+ KV<String, Integer> firstFoo = KV.of("foo", -1);
+ KV<String, Integer> secondFoo = KV.of("foo", 1);
+ KV<String, Integer> thirdFoo = KV.of("foo", 3);
+ KV<String, Integer> firstBar = KV.of("bar", 22);
+ KV<String, Integer> secondBar = KV.of("bar", 12);
+ KV<String, Integer> firstBaz = KV.of("baz", Integer.MAX_VALUE);
+ PCollection<KV<String, Integer>> values =
+ p.apply(Create.of(firstFoo, firstBar, secondFoo, firstBaz, secondBar, thirdFoo));
+ PCollection<KV<String, WindowedValue<Integer>>> kvs =
+ values.apply(new ReifyTimestampsAndWindows<String, Integer>());
+ PCollection<KeyedWorkItem<String, Integer>> groupedKvs =
+ kvs.apply(new DirectGroupByKeyOnly<String, Integer>());
+
+ CommittedBundle<KV<String, WindowedValue<Integer>>> inputBundle =
+ bundleFactory.createRootBundle(kvs).commit(Instant.now());
+ EvaluationContext evaluationContext = mock(EvaluationContext.class);
+
+ StructuralKey<String> fooKey = StructuralKey.of("foo", StringUtf8Coder.of());
+ UncommittedBundle<KeyedWorkItem<String, Integer>> fooBundle = bundleFactory.createKeyedBundle(
+ null, fooKey,
+ groupedKvs);
+ StructuralKey<String> barKey = StructuralKey.of("bar", StringUtf8Coder.of());
+ UncommittedBundle<KeyedWorkItem<String, Integer>> barBundle = bundleFactory.createKeyedBundle(
+ null, barKey,
+ groupedKvs);
+ StructuralKey<String> bazKey = StructuralKey.of("baz", StringUtf8Coder.of());
+ UncommittedBundle<KeyedWorkItem<String, Integer>> bazBundle = bundleFactory.createKeyedBundle(
+ null, bazKey,
+ groupedKvs);
+
+ when(evaluationContext.createKeyedBundle(inputBundle,
+ fooKey,
+ groupedKvs)).thenReturn(fooBundle);
+ when(evaluationContext.createKeyedBundle(inputBundle,
+ barKey,
+ groupedKvs)).thenReturn(barBundle);
+ when(evaluationContext.createKeyedBundle(inputBundle,
+ bazKey,
+ groupedKvs)).thenReturn(bazBundle);
+
+ // The input to a GroupByKey is assumed to be a KvCoder
+ @SuppressWarnings("unchecked")
+ Coder<String> keyCoder =
+ ((KvCoder<String, WindowedValue<Integer>>) kvs.getCoder()).getKeyCoder();
+ TransformEvaluator<KV<String, WindowedValue<Integer>>> evaluator =
+ new GroupByKeyOnlyEvaluatorFactory()
+ .forApplication(
+ groupedKvs.getProducingTransformInternal(), inputBundle, evaluationContext);
+
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(firstFoo)));
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(secondFoo)));
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(thirdFoo)));
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(firstBar)));
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(secondBar)));
+ evaluator.processElement(WindowedValue.valueInEmptyWindows(gwValue(firstBaz)));
+
+ evaluator.finishBundle();
+
+ assertThat(
+ fooBundle.commit(Instant.now()).getElements(),
+ contains(
+ new KeyedWorkItemMatcher<String, Integer>(
+ KeyedWorkItems.elementsWorkItem(
+ "foo",
+ ImmutableSet.of(
+ WindowedValue.valueInGlobalWindow(-1),
+ WindowedValue.valueInGlobalWindow(1),
+ WindowedValue.valueInGlobalWindow(3))),
+ keyCoder)));
+ assertThat(
+ barBundle.commit(Instant.now()).getElements(),
+ contains(
+ new KeyedWorkItemMatcher<String, Integer>(
+ KeyedWorkItems.elementsWorkItem(
+ "bar",
+ ImmutableSet.of(
+ WindowedValue.valueInGlobalWindow(12),
+ WindowedValue.valueInGlobalWindow(22))),
+ keyCoder)));
+ assertThat(
+ bazBundle.commit(Instant.now()).getElements(),
+ contains(
+ new KeyedWorkItemMatcher<String, Integer>(
+ KeyedWorkItems.elementsWorkItem(
+ "baz",
+ ImmutableSet.of(WindowedValue.valueInGlobalWindow(Integer.MAX_VALUE))),
+ keyCoder)));
+ }
+
+ private <K, V> KV<K, WindowedValue<V>> gwValue(KV<K, V> kv) {
+ return KV.of(kv.getKey(), WindowedValue.valueInGlobalWindow(kv.getValue()));
+ }
+
+ private static class KeyedWorkItemMatcher<K, V>
+ extends BaseMatcher<WindowedValue<KeyedWorkItem<K, V>>> {
+ private final KeyedWorkItem<K, V> myWorkItem;
+ private final Coder<K> keyCoder;
+
+ public KeyedWorkItemMatcher(KeyedWorkItem<K, V> myWorkItem, Coder<K> keyCoder) {
+ this.myWorkItem = myWorkItem;
+ this.keyCoder = keyCoder;
+ }
+
+ @Override
+ public boolean matches(Object item) {
+ if (item == null || !(item instanceof WindowedValue)) {
+ return false;
+ }
+ WindowedValue<KeyedWorkItem<K, V>> that = (WindowedValue<KeyedWorkItem<K, V>>) item;
+ Multiset<WindowedValue<V>> myValues = HashMultiset.create();
+ Multiset<WindowedValue<V>> thatValues = HashMultiset.create();
+ for (WindowedValue<V> value : myWorkItem.elementsIterable()) {
+ myValues.add(value);
+ }
+ for (WindowedValue<V> value : that.getValue().elementsIterable()) {
+ thatValues.add(value);
+ }
+ try {
+ return myValues.equals(thatValues)
+ && keyCoder
+ .structuralValue(myWorkItem.key())
+ .equals(keyCoder.structuralValue(that.getValue().key()));
+ } catch (Exception e) {
+ return false;
+ }
+ }
+
+ @Override
+ public void describeTo(Description description) {
+ description
+ .appendText("KeyedWorkItem<K, V> containing key ")
+ .appendValue(myWorkItem.key())
+ .appendText(" and values ")
+ .appendValueList("[", ", ", "]", myWorkItem.elementsIterable());
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactoryTest.java
index c4da86c..d40cf93 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactoryTest.java
@@ -57,7 +57,7 @@ public class ImmutabilityCheckingBundleFactoryTest {
TestPipeline p = TestPipeline.create();
created = p.apply(Create.<byte[]>of().withCoder(ByteArrayCoder.of()));
transformed = created.apply(ParDo.of(new IdentityDoFn<byte[]>()));
- factory = ImmutabilityCheckingBundleFactory.create(InProcessBundleFactory.create());
+ factory = ImmutabilityCheckingBundleFactory.create(ImmutableListBundleFactory.create());
}
@Test
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
index ead9c9e..890e06d 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
@@ -53,7 +53,7 @@ public class ImmutabilityEnforcementFactoryTest implements Serializable {
@Before
public void setup() {
factory = new ImmutabilityEnforcementFactory();
- bundleFactory = InProcessBundleFactory.create();
+ bundleFactory = ImmutableListBundleFactory.create();
TestPipeline p = TestPipeline.create();
pcollection =
p.apply(Create.of("foo".getBytes(), "spamhameggs".getBytes()))
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/4649eebe/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutableListBundleFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutableListBundleFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutableListBundleFactoryTest.java
new file mode 100644
index 0000000..21e4bcb
--- /dev/null
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutableListBundleFactoryTest.java
@@ -0,0 +1,231 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.direct;
+
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertThat;
+
+import org.apache.beam.runners.direct.DirectRunner.CommittedBundle;
+import org.apache.beam.runners.direct.DirectRunner.UncommittedBundle;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.coders.VoidCoder;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+
+import com.google.common.collect.ImmutableList;
+
+import org.hamcrest.Matcher;
+import org.hamcrest.Matchers;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+
+/**
+ * Tests for {@link ImmutableListBundleFactory}.
+ */
+@RunWith(JUnit4.class)
+public class ImmutableListBundleFactoryTest {
+ @Rule public ExpectedException thrown = ExpectedException.none();
+
+ private ImmutableListBundleFactory bundleFactory = ImmutableListBundleFactory.create();
+
+ private PCollection<Integer> created;
+ private PCollection<KV<String, Integer>> downstream;
+
+ @Before
+ public void setup() {
+ TestPipeline p = TestPipeline.create();
+ created = p.apply(Create.of(1, 2, 3));
+ downstream = created.apply(WithKeys.<String, Integer>of("foo"));
+ }
+
+ @Test
+ public void createRootBundleShouldCreateWithEmptyKey() {
+ PCollection<Integer> pcollection = TestPipeline.create().apply(Create.of(1));
+
+ UncommittedBundle<Integer> inFlightBundle = bundleFactory.createRootBundle(pcollection);
+
+ CommittedBundle<Integer> bundle = inFlightBundle.commit(Instant.now());
+
+ assertThat(bundle.getKey(),
+ Matchers.<StructuralKey<?>>equalTo(StructuralKey.of(null, VoidCoder.of())));
+ }
+
+ private <T> void createKeyedBundle(Coder<T> coder, T key) throws Exception {
+ PCollection<Integer> pcollection = TestPipeline.create().apply(Create.of(1));
+ StructuralKey skey = StructuralKey.of(key, coder);
+
+ UncommittedBundle<Integer> inFlightBundle =
+ bundleFactory.createKeyedBundle(null, skey, pcollection);
+
+ CommittedBundle<Integer> bundle = inFlightBundle.commit(Instant.now());
+ assertThat(bundle.getKey(), equalTo(skey));
+ }
+
+ @Test
+ public void keyedWithNullKeyShouldCreateKeyedBundle() throws Exception {
+ createKeyedBundle(VoidCoder.of(), null);
+ }
+
+ @Test
+ public void keyedWithKeyShouldCreateKeyedBundle() throws Exception {
+ createKeyedBundle(StringUtf8Coder.of(), "foo");
+ createKeyedBundle(VarIntCoder.of(), 1234);
+ createKeyedBundle(ByteArrayCoder.of(), new byte[] {0, 2, 4, 99});
+ }
+
+ private <T> CommittedBundle<T>
+ afterCommitGetElementsShouldHaveAddedElements(Iterable<WindowedValue<T>> elems) {
+ PCollection<T> pcollection = TestPipeline.create().apply(Create.<T>of());
+
+ UncommittedBundle<T> bundle = bundleFactory.createRootBundle(pcollection);
+ Collection<Matcher<? super WindowedValue<T>>> expectations = new ArrayList<>();
+ for (WindowedValue<T> elem : elems) {
+ bundle.add(elem);
+ expectations.add(equalTo(elem));
+ }
+ Matcher<Iterable<? extends WindowedValue<T>>> containsMatcher =
+ Matchers.<WindowedValue<T>>containsInAnyOrder(expectations);
+ CommittedBundle<T> committed = bundle.commit(Instant.now());
+ assertThat(committed.getElements(), containsMatcher);
+
+ return committed;
+ }
+
+ @Test
+ public void getElementsBeforeAddShouldReturnEmptyIterable() {
+ afterCommitGetElementsShouldHaveAddedElements(Collections.<WindowedValue<Integer>>emptyList());
+ }
+
+ @Test
+ public void getElementsAfterAddShouldReturnAddedElements() {
+ WindowedValue<Integer> firstValue = WindowedValue.valueInGlobalWindow(1);
+ WindowedValue<Integer> secondValue =
+ WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1000L));
+
+ afterCommitGetElementsShouldHaveAddedElements(ImmutableList.of(firstValue, secondValue));
+ }
+
+ @SuppressWarnings("unchecked")
+ @Test
+ public void withElementsShouldReturnIndependentBundle() {
+ WindowedValue<Integer> firstValue = WindowedValue.valueInGlobalWindow(1);
+ WindowedValue<Integer> secondValue =
+ WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1000L));
+
+ CommittedBundle<Integer> committed =
+ afterCommitGetElementsShouldHaveAddedElements(ImmutableList.of(firstValue, secondValue));
+
+ WindowedValue<Integer> firstReplacement =
+ WindowedValue.of(
+ 9,
+ new Instant(2048L),
+ new IntervalWindow(new Instant(2044L), Instant.now()),
+ PaneInfo.NO_FIRING);
+ WindowedValue<Integer> secondReplacement =
+ WindowedValue.timestampedValueInGlobalWindow(-1, Instant.now());
+ CommittedBundle<Integer> withed =
+ committed.withElements(ImmutableList.of(firstReplacement, secondReplacement));
+
+ assertThat(withed.getElements(), containsInAnyOrder(firstReplacement, secondReplacement));
+ assertThat(committed.getElements(), containsInAnyOrder(firstValue, secondValue));
+ assertThat(withed.getKey(), Matchers.<StructuralKey<?>>equalTo(committed.getKey()));
+ assertThat(withed.getPCollection(), equalTo(committed.getPCollection()));
+ assertThat(
+ withed.getSynchronizedProcessingOutputWatermark(),
+ equalTo(committed.getSynchronizedProcessingOutputWatermark()));
+ }
+
+ @Test
+ public void addAfterCommitShouldThrowException() {
+ PCollection<Integer> pcollection = TestPipeline.create().apply(Create.<Integer>of());
+
+ UncommittedBundle<Integer> bundle = bundleFactory.createRootBundle(pcollection);
+ bundle.add(WindowedValue.valueInGlobalWindow(1));
+ CommittedBundle<Integer> firstCommit = bundle.commit(Instant.now());
+ assertThat(firstCommit.getElements(), containsInAnyOrder(WindowedValue.valueInGlobalWindow(1)));
+
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("3");
+ thrown.expectMessage("committed");
+
+ bundle.add(WindowedValue.valueInGlobalWindow(3));
+ }
+
+ @Test
+ public void commitAfterCommitShouldThrowException() {
+ PCollection<Integer> pcollection = TestPipeline.create().apply(Create.<Integer>of());
+
+ UncommittedBundle<Integer> bundle = bundleFactory.createRootBundle(pcollection);
+ bundle.add(WindowedValue.valueInGlobalWindow(1));
+ CommittedBundle<Integer> firstCommit = bundle.commit(Instant.now());
+ assertThat(firstCommit.getElements(), containsInAnyOrder(WindowedValue.valueInGlobalWindow(1)));
+
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("committed");
+
+ bundle.commit(Instant.now());
+ }
+
+ @Test
+ public void createBundleUnkeyedResultUnkeyed() {
+ CommittedBundle<KV<String, Integer>> newBundle =
+ bundleFactory
+ .createBundle(bundleFactory.createRootBundle(created).commit(Instant.now()), downstream)
+ .commit(Instant.now());
+ }
+
+ @Test
+ public void createBundleKeyedResultPropagatesKey() {
+ CommittedBundle<KV<String, Integer>> newBundle =
+ bundleFactory.createBundle(
+ bundleFactory.createKeyedBundle(
+ null,
+ StructuralKey.of("foo", StringUtf8Coder.of()),
+ created).commit(Instant.now()),
+ downstream).commit(Instant.now());
+ assertThat(newBundle.getKey().getKey(), Matchers.<Object>equalTo("foo"));
+ }
+
+ @Test
+ public void createKeyedBundleKeyed() {
+ CommittedBundle<KV<String, Integer>> keyedBundle = bundleFactory.createKeyedBundle(
+ bundleFactory.createRootBundle(created).commit(Instant.now()),
+ StructuralKey.of("foo", StringUtf8Coder.of()),
+ downstream).commit(Instant.now());
+ assertThat(keyedBundle.getKey().getKey(), Matchers.<Object>equalTo("foo"));
+ }
+}