Examples with InstructionResponse org.apache.beam.fn.v1.BeamFnApi.InstructionResponse used on opensource projects

Example 1 with InstructionResponse

use of org.apache.beam.fn.v1.BeamFnApi.InstructionResponse in project beam by apache.

the class FnHarnessTest method testLaunchFnHarnessAndTeardownCleanly.

@Test
public void testLaunchFnHarnessAndTeardownCleanly() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    List<BeamFnApi.LogEntry> logEntries = new ArrayList<>();
    List<BeamFnApi.InstructionResponse> instructionResponses = new ArrayList<>();
    BeamFnLoggingGrpc.BeamFnLoggingImplBase loggingService = new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<LogControl> responseObserver) {
            return TestStreams.withOnNext((BeamFnApi.LogEntry.List entries) -> logEntries.addAll(entries.getLogEntriesList())).withOnCompleted(() -> responseObserver.onCompleted()).build();
        }
    };
    BeamFnControlGrpc.BeamFnControlImplBase controlService = new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<InstructionResponse> control(StreamObserver<InstructionRequest> responseObserver) {
            CountDownLatch waitForResponses = new CountDownLatch(1);
            options.as(GcsOptions.class).getExecutorService().submit(new Runnable() {

                @Override
                public void run() {
                    responseObserver.onNext(INSTRUCTION_REQUEST);
                    Uninterruptibles.awaitUninterruptibly(waitForResponses);
                    responseObserver.onCompleted();
                }
            });
            return TestStreams.withOnNext(new Consumer<BeamFnApi.InstructionResponse>() {

                @Override
                public void accept(InstructionResponse t) {
                    instructionResponses.add(t);
                    waitForResponses.countDown();
                }
            }).withOnCompleted(waitForResponses::countDown).build();
        }
    };
    Server loggingServer = ServerBuilder.forPort(0).addService(loggingService).build();
    loggingServer.start();
    try {
        Server controlServer = ServerBuilder.forPort(0).addService(controlService).build();
        controlServer.start();
        try {
            BeamFnApi.ApiServiceDescriptor loggingDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setId("1L").setUrl("localhost:" + loggingServer.getPort()).build();
            BeamFnApi.ApiServiceDescriptor controlDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setId("2L").setUrl("localhost:" + controlServer.getPort()).build();
            FnHarness.main(options, loggingDescriptor, controlDescriptor);
            assertThat(instructionResponses, contains(INSTRUCTION_RESPONSE));
        } finally {
            controlServer.shutdownNow();
        }
    } finally {
        loggingServer.shutdownNow();
    }
}