BeamWorkerStatusGrpcService extends BeamFnWorkerStatusImplBase implements FnService { public Map<String, String> getAllWorkerStatuses(long timeout, TimeUnit timeUnit) { if (isClosed.get()) { throw new IllegalStateException("BeamWorkerStatusGrpcService already closed."); } Map<String, String> allStatuses = new ConcurrentSkipListMap<>(Comparator.naturalOrder()); Set<String> connectedClientIdsCopy; synchronized (connectedClient) { connectedClientIdsCopy = ImmutableSet.copyOf(connectedClient.keySet()); } connectedClientIdsCopy .parallelStream() .forEach( workerId -> allStatuses.put(workerId, getSingleWorkerStatus(workerId, timeout, timeUnit))); return allStatuses; } private  BeamWorkerStatusGrpcService(
      ApiServiceDescriptor apiServiceDescriptor, HeaderAccessor headerAccessor); static BeamWorkerStatusGrpcService create(
      ApiServiceDescriptor apiServiceDescriptor, HeaderAccessor headerAccessor); @Override void close(); @Override StreamObserver<WorkerStatusResponse> workerStatus(
      StreamObserver<WorkerStatusRequest> requestObserver); String getSingleWorkerStatus(String workerId, long timeout, TimeUnit timeUnit); Map<String, String> getAllWorkerStatuses(long timeout, TimeUnit timeUnit);  }
@Test public void testGetAllWorkerStatuses() throws Exception { Set<String> ids = Sets.newHashSet("id0", "id3", "id11", "id12", "id21"); for (String id : ids) { StreamObserver<WorkerStatusRequest> requestObserverMock = mock(StreamObserver.class); BeamFnWorkerStatusStub workerStatusStub = BeamFnWorkerStatusGrpc.newStub(channel) .withInterceptors(AddHarnessIdInterceptor.create(id)); StreamObserver<WorkerStatusResponse> observer = workerStatusStub.workerStatus(requestObserverMock); waitAndGetStatusClient(id); doAnswer( (invocation) -> { WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0]; observer.onNext( WorkerStatusResponse.newBuilder() .setId(request.getId()) .setStatusInfo("status") .build()); return null; }) .when(requestObserverMock) .onNext(any()); } Map<String, String> allWorkerStatuses = service.getAllWorkerStatuses(5, TimeUnit.SECONDS); assertEquals(ids, allWorkerStatuses.keySet()); for (String id : ids) { assertEquals("status", allWorkerStatuses.get(id)); } }