public Optional<ResT> waitForComplete(long timeoutMs) throws IOException {
    if (mCompleted || mCanceled) {
      return Optional.empty();
    }
    ResT prevResponse;
    ResT response = null;
    do {
      // wait until inbound stream is closed from server.
      prevResponse = response;
      response = receive(timeoutMs);
    } while (response != null);
    return Optional.ofNullable(prevResponse);
  }
@Test
  public void waitForComplete() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE * 2).toArray(WriteResponse[]::new);

    EXECUTOR.submit(() -> {
      for (WriteResponse response : responses) {
        mResponseObserver.onNext(response);
      }
      try {
        Thread.sleep(SHORT_TIMEOUT);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
      mResponseObserver.onCompleted();
    });

    WriteResponse actualResponse = mStream.receive(TIMEOUT);
    assertEquals(responses[0], actualResponse);

    mStream.waitForComplete(TIMEOUT);

    actualResponse = mStream.receive(TIMEOUT);
    assertEquals(null, actualResponse);
  }