use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.
the class SpannerIOWriteTest method testBatchableMutationFilterFn_size.
@Test
public void testBatchableMutationFilterFn_size() {
Mutation all = Mutation.delete("test", KeySet.all());
Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too big.
g(m(1L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range) };
long mutationSize = MutationSizeEstimator.sizeOf(m(1L));
BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, mutationSize * 3, 1000, 1000);
DoFn<MutationGroup, MutationGroup>.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
// Capture the outputs.
doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
// Process all elements.
for (MutationGroup m : mutationGroups) {
when(mockProcessContext.element()).thenReturn(m);
testFn.processElement(mockProcessContext);
}
// Verify captured batchable elements.
assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
// Verify captured unbatchable mutations
Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too big.
g(m(1L), m(3L), m(4L), m(5L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range)));
}
use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.
the class DoFnInvokersTest method testSplittableDoFnWithAllMethods.
@Test
public void testSplittableDoFnWithAllMethods() throws Exception {
MockFn fn = mock(MockFn.class);
DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
final SomeRestrictionTracker tracker = mock(SomeRestrictionTracker.class);
final SomeRestrictionCoder coder = mock(SomeRestrictionCoder.class);
final InstantCoder watermarkEstimatorStateCoder = InstantCoder.of();
final Instant watermarkEstimatorState = Instant.now();
final WatermarkEstimator<Instant> watermarkEstimator = new WatermarkEstimators.Manual(watermarkEstimatorState);
SomeRestriction restriction = new SomeRestriction();
final SomeRestriction part1 = new SomeRestriction();
final SomeRestriction part2 = new SomeRestriction();
final SomeRestriction part3 = new SomeRestriction();
when(fn.getRestrictionCoder()).thenReturn(coder);
when(fn.getWatermarkEstimatorStateCoder()).thenReturn(watermarkEstimatorStateCoder);
when(fn.getInitialRestriction(mockElement)).thenReturn(restriction);
doAnswer(AdditionalAnswers.delegatesTo(new MockFn() {
@DoFn.SplitRestriction
@Override
public void splitRestriction(@Element String element, @Restriction SomeRestriction restriction, DoFn.OutputReceiver<SomeRestriction> receiver) {
receiver.output(part1);
receiver.output(part2);
receiver.output(part3);
}
})).when(fn).splitRestriction(eq(mockElement), same(restriction), any());
when(fn.getInitialWatermarkEstimatorState()).thenReturn(watermarkEstimatorState);
when(fn.newTracker(restriction)).thenReturn(tracker);
when(fn.newWatermarkEstimator(watermarkEstimatorState)).thenReturn(watermarkEstimator);
when(fn.processElement(mockProcessContext, tracker, watermarkEstimator)).thenReturn(resume());
when(fn.getSize()).thenReturn(2.0);
assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault()));
assertEquals(watermarkEstimatorStateCoder, invoker.invokeGetWatermarkEstimatorStateCoder(CoderRegistry.createDefault()));
assertEquals(restriction, invoker.invokeGetInitialRestriction(new FakeArgumentProvider<String, String>() {
@Override
public String element(DoFn<String, String> doFn) {
return mockElement;
}
}));
List<SomeRestriction> outputs = new ArrayList<>();
invoker.invokeSplitRestriction(new FakeArgumentProvider<String, String>() {
@Override
public String element(DoFn<String, String> doFn) {
return mockElement;
}
@Override
public Object restriction() {
return restriction;
}
@Override
public OutputReceiver outputReceiver(DoFn doFn) {
return new OutputReceiver<SomeRestriction>() {
@Override
public void output(SomeRestriction output) {
outputs.add(output);
}
@Override
public void outputWithTimestamp(SomeRestriction output, Instant timestamp) {
fail("Unexpected output with timestamp");
}
};
}
});
assertEquals(Arrays.asList(part1, part2, part3), outputs);
assertEquals(watermarkEstimatorState, invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>()));
assertEquals(tracker, invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {
@Override
public String element(DoFn<String, String> doFn) {
return mockElement;
}
@Override
public Object restriction() {
return restriction;
}
}));
assertEquals(watermarkEstimator, invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {
@Override
public Object watermarkEstimatorState() {
return watermarkEstimatorState;
}
}));
assertEquals(resume(), invoker.invokeProcessElement(new FakeArgumentProvider<String, String>() {
@Override
public DoFn<String, String>.ProcessContext processContext(DoFn<String, String> fn) {
return mockProcessContext;
}
@Override
public RestrictionTracker<?, ?> restrictionTracker() {
return tracker;
}
@Override
public WatermarkEstimator<?> watermarkEstimator() {
return watermarkEstimator;
}
}));
assertEquals(2.0, invoker.invokeGetSize(mockArgumentProvider), 0.0001);
}
use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.
the class DoFnInvokersTest method testTruncateFnWithHasDefaultMethodsWhenBounded.
@Test
public void testTruncateFnWithHasDefaultMethodsWhenBounded() throws Exception {
class BoundedMockFn extends DoFn<String, String> {
@ProcessElement
public void processElement(ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {
}
@GetInitialRestriction
public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) {
return null;
}
@GetInitialWatermarkEstimatorState
public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() {
return null;
}
}
BoundedMockFn fn = mock(BoundedMockFn.class);
DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
CoderRegistry coderRegistry = CoderRegistry.createDefault();
coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class));
coderRegistry.registerCoderForClass(WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
assertThat(invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
assertThat(invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
RestrictionTracker tracker = invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {
@Override
public Object restriction() {
return new RestrictionWithBoundedDefaultTracker();
}
});
assertThat(tracker, instanceOf(BoundedDefaultTracker.class));
TruncateResult<?> result = invoker.invokeTruncateRestriction(new FakeArgumentProvider<String, String>() {
@Override
public RestrictionTracker restrictionTracker() {
return tracker;
}
@Override
public String element(DoFn<String, String> doFn) {
return "blah";
}
@Override
public Object restriction() {
return "foo";
}
});
assertEquals("foo", result.getTruncatedRestriction());
assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
assertThat(invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {
@Override
public Object watermarkEstimatorState() {
return new WatermarkEstimatorStateWithDefaultWatermarkEstimator();
}
}), instanceOf(DefaultWatermarkEstimator.class));
}
use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.
the class DoFnSignaturesSplittableDoFnTest method testHasRestrictionTracker.
@Test
// used via reflection
@SuppressWarnings("unused")
public void testHasRestrictionTracker() throws Exception {
DoFnSignature.ProcessElementMethod signature = analyzeProcessElementMethod(new AnonymousMethod() {
private void method(DoFn<Integer, String>.ProcessContext context, SomeRestrictionTracker tracker) {
}
});
assertTrue(signature.isSplittable());
assertTrue(signature.extraParameters().stream().anyMatch(Predicates.instanceOf(DoFnSignature.Parameter.RestrictionTrackerParameter.class)::apply));
assertEquals(SomeRestrictionTracker.class, signature.trackerT().getRawType());
}
use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.
the class DoFnInvokersTest method testDoFnWithTimer.
/**
* Tests that the generated {@link DoFnInvoker} passes the timer parameter that it should.
*/
@Test
public void testDoFnWithTimer() throws Exception {
Timer mockTimer = mock(Timer.class);
final String timerId = "my-timer-id-here";
when(mockArgumentProvider.timer(TimerDeclaration.PREFIX + timerId)).thenReturn(mockTimer);
class MockFn extends DoFn<String, String> {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@ProcessElement
public void processElement(ProcessContext c, @TimerId(timerId) Timer timer) throws Exception {
}
@OnTimer(timerId)
public void onTimer() {
}
}
MockFn fn = mock(MockFn.class);
assertEquals(stop(), invokeProcessElement(fn));
verify(fn).processElement(mockProcessContext, mockTimer);
}
Aggregations