http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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/babddbbc/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")); + } +}
