use of org.apache.beam.runners.dataflow.worker.util.common.worker.TestOutputReceiver in project beam by apache.
the class PartialGroupByKeyParDoFnsTest method testPartialGroupByKey.
@Test
public void testPartialGroupByKey() throws Exception {
Coder keyCoder = StringUtf8Coder.of();
Coder valueCoder = BigEndianIntegerCoder.of();
TestOutputReceiver receiver = new TestOutputReceiver(new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(KvCoder.of(keyCoder, IterableCoder.of(valueCoder)))), counterSet, NameContextsForTests.nameContextForTest());
ParDoFn pgbkParDoFn = new SimplePartialGroupByKeyParDoFn(GroupingTables.buffering(new WindowingCoderGroupingKeyCreator(keyCoder), PairInfo.create(), new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)), new CoderSizeEstimator(valueCoder)), receiver);
pgbkParDoFn.startBundle(receiver);
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 4)));
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 5)));
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 6)));
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("joe", 7)));
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 8)));
pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 9)));
pgbkParDoFn.finishBundle();
assertThat(receiver.outputElems, IsIterableContainingInAnyOrder.<Object>containsInAnyOrder(WindowedValue.valueInGlobalWindow(KV.of("hi", Arrays.asList(4, 6, 9))), WindowedValue.valueInGlobalWindow(KV.of("there", Arrays.asList(5, 8))), WindowedValue.valueInGlobalWindow(KV.of("joe", Arrays.asList(7)))));
// Exact counter values depend on size of encoded data. If encoding
// changes, then these expected counters should change to match.
CounterUpdateExtractor<?> updateExtractor = Mockito.mock(CounterUpdateExtractor.class);
counterSet.extractUpdates(false, updateExtractor);
verify(updateExtractor).longSum(getObjectCounterName("test_receiver_out"), false, 3L);
verify(updateExtractor).longMean(getMeanByteCounterName("test_receiver_out"), false, LongCounterMean.ZERO.addValue(49L, 3));
verifyNoMoreInteractions(updateExtractor);
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.TestOutputReceiver in project beam by apache.
the class PartialGroupByKeyParDoFnsTest method testCreateWithCombinerAndStreamingSideInputs.
@Test
public void testCreateWithCombinerAndStreamingSideInputs() throws Exception {
StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class);
options.setStreaming(true);
Coder keyCoder = StringUtf8Coder.of();
Coder valueCoder = BigEndianIntegerCoder.of();
KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);
TestOutputReceiver receiver = new TestOutputReceiver(new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)), counterSet, NameContextsForTests.nameContextForTest());
when(mockSideInputReader.isEmpty()).thenReturn(false);
when(mockStreamingStepContext.stateInternals()).thenReturn((StateInternals) mockStateInternals);
when(mockStateInternals.state(Matchers.<StateNamespace>any(), Matchers.<StateTag>any())).thenReturn(mockState);
when(mockState.read()).thenReturn(Maps.newHashMap());
ParDoFn pgbk = PartialGroupByKeyParDoFns.create(options, kvCoder, AppliedCombineFn.withInputCoder(Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder, ImmutableList.<PCollectionView<?>>of(), WindowingStrategy.globalDefault()), mockSideInputReader, receiver, mockStreamingStepContext);
assertTrue(pgbk instanceof StreamingSideInputPGBKParDoFn);
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.TestOutputReceiver in project beam by apache.
the class IntrinsicMapTaskExecutorTest method testValidOperations.
@Test
public void testValidOperations() throws Exception {
TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest());
List<Operation> operations = Arrays.<Operation>asList(new TestReadOperation(receiver, createContext("ReadOperation")));
ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
Assert.assertEquals(operations.get(0), executor.getReadOperation());
}
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.TestOutputReceiver in project beam by apache.
the class RemoteGrpcPortReadOperationTest method setUp.
@Before
public void setUp() {
MockitoAnnotations.initMocks(this);
testReceiver = new TestOutputReceiver(CODER, NameContextsForTests.nameContextForTest());
operation = new RemoteGrpcPortReadOperation<>(beamFnDataService, TRANSFORM_ID, bundleIdSupplier, CODER, new OutputReceiver[] { testReceiver }, operationContext);
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.TestOutputReceiver in project beam by apache.
the class IntrinsicMapTaskExecutorTest method testGetProgressAndRequestSplit.
@Test
public void testGetProgressAndRequestSplit() throws Exception {
TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest());
TestReadOperation operation = new TestReadOperation(receiver, createContext("ReadOperation"));
ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(Arrays.asList(new Operation[] { operation }), counterSet, stateTracker)) {
operation.setProgress(approximateProgressAtIndex(1L));
Assert.assertEquals(positionAtIndex(1L), positionFromProgress(executor.getWorkerProgress()));
Assert.assertEquals(positionAtIndex(1L), positionFromSplitResult(executor.requestDynamicSplit(splitRequestAtIndex(1L))));
}
}
Aggregations