FnHarness { public static void main(String[] args) throws Exception { main(System::getenv); }  static void main(String[] args); @VisibleForTesting static void main(Function<String, String> environmentVarGetter); static void main(
      String id,
      PipelineOptions options,
      Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor,
      Endpoints.ApiServiceDescriptor controlApiServiceDescriptor); static void main(
      String id,
      PipelineOptions options,
      Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor,
      Endpoints.ApiServiceDescriptor controlApiServiceDescriptor,
      ManagedChannelFactory channelFactory,
      OutboundObserverFactory outboundObserverFactory);  }
@Test @SuppressWarnings("FutureReturnValueIgnored") public void testLaunchFnHarnessAndTeardownCleanly() throws Exception { Function<String, String> environmentVariableMock = mock(Function.class); PipelineOptions options = PipelineOptionsFactory.create(); when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id"); when(environmentVariableMock.apply("PIPELINE_OPTIONS")) .thenReturn(PipelineOptionsTranslation.toJson(options)); List<BeamFnApi.LogEntry> logEntries = new ArrayList<>(); List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class); 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( () -> { responseObserver.onNext(INSTRUCTION_REQUEST); Uninterruptibles.awaitUninterruptibly(waitForResponses); responseObserver.onCompleted(); }); return TestStreams.withOnNext( (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 { Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl("localhost:" + loggingServer.getPort()) .build(); Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl("localhost:" + controlServer.getPort()) .build(); when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")) .thenReturn(TextFormat.printToString(loggingDescriptor)); when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")) .thenReturn(TextFormat.printToString(controlDescriptor)); FnHarness.main(environmentVariableMock); } finally { controlServer.shutdownNow(); } } finally { loggingServer.shutdownNow(); } InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses); inOrder.verify(onStartupMock).run(); inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any()); inOrder.verify(beforeProcessingMock).accept(any()); inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE); }
@Test(timeout = 10 * 1000) @SuppressWarnings("FutureReturnValueIgnored") public void testLaunchFnHarnessAndTeardownCleanly() throws Exception { Function<String, String> environmentVariableMock = mock(Function.class); PipelineOptions options = PipelineOptionsFactory.create(); when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id"); when(environmentVariableMock.apply("PIPELINE_OPTIONS")) .thenReturn(PipelineOptionsTranslation.toJson(options)); List<BeamFnApi.LogEntry> logEntries = new ArrayList<>(); List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class); 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( () -> { responseObserver.onNext(INSTRUCTION_REQUEST); Uninterruptibles.awaitUninterruptibly(waitForResponses); responseObserver.onCompleted(); }); return TestStreams.withOnNext( (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 { Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl("localhost:" + loggingServer.getPort()) .build(); Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl("localhost:" + controlServer.getPort()) .build(); when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")) .thenReturn(TextFormat.printToString(loggingDescriptor)); when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")) .thenReturn(TextFormat.printToString(controlDescriptor)); FnHarness.main(environmentVariableMock); } finally { controlServer.shutdownNow(); } } finally { loggingServer.shutdownNow(); } InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses); inOrder.verify(onStartupMock).run(); inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any()); inOrder.verify(beforeProcessingMock).accept(any()); inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE); }