public Map<String, String> getAllWorkerStatuses(long timeout, TimeUnit timeUnit) {
    if (isClosed.get()) {
      throw new IllegalStateException("BeamWorkerStatusGrpcService already closed.");
    }
    // return result in worker id sorted map.
    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;
  }
@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);
      // wait for connection before proceeding to avoid race condition.
      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));
    }
  }