@Override
  public GetDataStream getDataStream() {
    return windmillStreamFactory.createGetDataStream(
        dispatcherClient.getWindmillServiceStub(), throttleTimers.getDataThrottleTimer());
  }
@Test
  public void testStreamingGetDataHeartbeatsAsHeartbeatRequests() throws Exception {
    // This server records the heartbeats observed but doesn't respond.
    final List<ComputationHeartbeatRequest> receivedHeartbeats = new ArrayList<>();

    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetDataRequest> getDataStream(
              StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetDataRequest chunk) {
                try {
                  if (!sawHeader) {
                    LOG.info("Received header");
                    errorCollector.checkThat(
                        chunk.getHeader(),
                        Matchers.equalTo(
                            JobHeader.newBuilder()
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setClientId(clientId)
                                .build()));
                    sawHeader = true;
                  } else {
                    LOG.info(
                        "Received {} computationHeartbeatRequests",
                        chunk.getComputationHeartbeatRequestCount());
                    errorCollector.checkThat(
                        chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));
                    errorCollector.checkThat(chunk.getRequestIdCount(), Matchers.is(0));

                    synchronized (receivedHeartbeats) {
                      receivedHeartbeats.addAll(chunk.getComputationHeartbeatRequestList());
                    }
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                responseObserver.onCompleted();
              }
            };
          }
        });

    List<String> computation1Keys = new ArrayList<>();
    List<String> computation2Keys = new ArrayList<>();

    // When sending heartbeats as HeartbeatRequest protos, all keys for the same computation should
    // be batched into the same ComputationHeartbeatRequest. Compare to the KeyedGetDataRequest
    // version in the test above, which only sends one key per ComputationGetDataRequest.
    List<ComputationHeartbeatRequest> expectedHeartbeats = new ArrayList<>();
    ComputationHeartbeatRequest.Builder comp1Builder =
        ComputationHeartbeatRequest.newBuilder().setComputationId("Computation1");
    ComputationHeartbeatRequest.Builder comp2Builder =
        ComputationHeartbeatRequest.newBuilder().setComputationId("Computation2");
    for (int i = 0; i < 100; ++i) {
      String computation1Key = "Computation1Key" + i;
      computation1Keys.add(computation1Key);
      comp1Builder.addHeartbeatRequests(
          makeHeartbeatRequest(Collections.singletonList(computation1Key)).get(0));
      String computation2Key = "Computation2Key" + largeString(i * 20);
      computation2Keys.add(computation2Key);
      comp2Builder.addHeartbeatRequests(
          makeHeartbeatRequest(Collections.singletonList(computation2Key)).get(0));
    }
    expectedHeartbeats.add(comp1Builder.build());
    expectedHeartbeats.add(comp2Builder.build());
    Map<String, Collection<HeartbeatRequest>> heartbeatRequestMap = new HashMap<>();
    heartbeatRequestMap.put("Computation1", makeHeartbeatRequest(computation1Keys));
    heartbeatRequestMap.put("Computation2", makeHeartbeatRequest(computation2Keys));

    GetDataStream stream = client.getDataStream();
    stream.refreshActiveWork(heartbeatRequestMap);
    stream.halfClose();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));

    boolean receivedAllHeartbeatRequests = false;
    while (!receivedAllHeartbeatRequests) {
      Thread.sleep(100);
      synchronized (receivedHeartbeats) {
        if (receivedHeartbeats.size() != expectedHeartbeats.size()) {
          continue;
        }
        assertEquals(expectedHeartbeats, receivedHeartbeats);
        receivedAllHeartbeatRequests = true;
      }
    }
  }