Examples with FnDataReceiver org.apache.beam.sdk.fn.data.FnDataReceiver used on opensource projects

Example 1 with FnDataReceiver

use of org.apache.beam.sdk.fn.data.FnDataReceiver in project flink by apache.

the class PythonTestUtils method createMockJobBundleFactory.

public static JobBundleFactory createMockJobBundleFactory() {
    JobBundleFactory jobBundleFactorySpy = spy(JobBundleFactory.class);
    StageBundleFactory stageBundleFactorySpy = spy(StageBundleFactory.class);
    when(jobBundleFactorySpy.forStage(any())).thenReturn(stageBundleFactorySpy);
    ProcessBundleDescriptors.ExecutableProcessBundleDescriptor processBundleDescriptor = spy(ProcessBundleDescriptors.ExecutableProcessBundleDescriptor.class);
    when(stageBundleFactorySpy.getProcessBundleDescriptor()).thenReturn(processBundleDescriptor);
    RemoteBundle remoteBundleSpy = spy(RemoteBundle.class);
    try {
        when(stageBundleFactorySpy.getBundle(any(OutputReceiverFactory.class), any(TimerReceiverFactory.class), any(StateRequestHandler.class), any(BundleProgressHandler.class))).thenReturn(remoteBundleSpy);
    } catch (Exception e) {
    // ignore
    }
    Map<String, FnDataReceiver> inputReceivers = new HashMap<>();
    FnDataReceiver<WindowedValue<?>> windowedValueReceiverSpy = spy(FnDataReceiver.class);
    inputReceivers.put("input", windowedValueReceiverSpy);
    when(remoteBundleSpy.getInputReceivers()).thenReturn(inputReceivers);
    return jobBundleFactorySpy;
}