Examples with PTransformRunnerFactory org.apache.beam.fn.harness.PTransformRunnerFactory used on opensource projects

Example 11 with PTransformRunnerFactory

use of org.apache.beam.fn.harness.PTransformRunnerFactory in project beam by apache.

the class ProcessBundleHandlerTest method testBundleFinalizationIsPropagated.

@Test
public void testBundleFinalizationIsPropagated() throws Exception {
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    FinalizeBundleHandler mockFinalizeBundleHandler = mock(FinalizeBundleHandler.class);
    BundleFinalizer.Callback mockCallback = mock(BundleFinalizer.Callback.class);
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
    mockFinalizeBundleHandler, new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (context) -> {
        BundleFinalizer bundleFinalizer = context.getBundleFinalizer();
        context.addStartBundleFunction(() -> bundleFinalizer.afterBundleCommit(Instant.ofEpochMilli(42L), mockCallback));
        return null;
    }), Caches.noop(), new BundleProcessorCache());
    BeamFnApi.InstructionResponse.Builder response = handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("2L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    assertTrue(response.getProcessBundle().getRequiresFinalization());
    verify(mockFinalizeBundleHandler).registerCallbacks(eq("2L"), argThat((Collection<CallbackRegistration> arg) -> {
        CallbackRegistration registration = Iterables.getOnlyElement(arg);
        assertEquals(Instant.ofEpochMilli(42L), registration.getExpiryTime());
        assertSame(mockCallback, registration.getCallback());
        return true;
    }));
}