Code review comment: Instead of just sleeping could you poll some counter you increment when returning server errors, and wait for a couple errors to be returned?  That would help ensure the error handling path on the client actually runs.

Code before review:
public void testStreamingCommitClosedStream() throws Exception {
List<WorkItemCommitRequest> commitRequestList = new ArrayList<>();
List<CountDownLatch> latches = new ArrayList<>();
Map<Long, WorkItemCommitRequest> commitRequests = new ConcurrentHashMap<>();
AtomicBoolean shouldServerReturnError = new AtomicBoolean(true);
for (int i = 0; i < 500; ++i) {
WorkItemCommitRequest request = makeCommitRequest(i, i * (i < 480 ? 8 : 128));
commitRequestList.add(request);
commitRequests.put((long) i, request);
latches.add(new CountDownLatch(1));
}
Collections.shuffle(commitRequestList);
serviceRegistry.addService(
new CloudWindmillServiceV1Alpha1ImplBase() {
@Override
public StreamObserver<StreamingCommitWorkRequest> commitWorkStream(
StreamObserver<StreamingCommitResponse> responseObserver) {
StreamObserver<StreamingCommitWorkRequest> testCommitStreamObserver =
getTestCommitStreamObserver(responseObserver, commitRequests);
return new StreamObserver<StreamingCommitWorkRequest>() {
@Override
public void onNext(StreamingCommitWorkRequest request) {
if (shouldServerReturnError.get()) {
try {
responseObserver.onError(
new RuntimeException("shouldServerReturnError = true"));
} catch (IllegalStateException e) {
}
} else {
testCommitStreamObserver.onNext(request);
}
}
@Override
public void onError(Throwable throwable) {
testCommitStreamObserver.onError(throwable);
}
@Override
public void onCompleted() {
testCommitStreamObserver.onCompleted();
}
};
}
});
CommitWorkStream stream = client.commitWorkStream();
for (int i = 0; i < commitRequestList.size(); ) {
final CountDownLatch latch = latches.get(i);
if (stream.commitWorkItem(
"computation",
commitRequestList.get(i),
(CommitStatus status) -> {
assertEquals(status, CommitStatus.OK);
latch.countDown();
})) {
i++;
} else {
stream.flush();
}
}
stream.flush();
stream.close();
Thread.sleep(100);
shouldServerReturnError.set(false);
for (CountDownLatch latch : latches) {
assertTrue(latch.await(1, TimeUnit.MINUTES));
}
assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
}

Line under review: Thread.sleep(100);

Code after addressing the review:
public void testStreamingCommitClosedStream() throws Exception {
List<WorkItemCommitRequest> commitRequestList = new ArrayList<>();
List<CountDownLatch> latches = new ArrayList<>();
Map<Long, WorkItemCommitRequest> commitRequests = new ConcurrentHashMap<>();
AtomicBoolean shouldServerReturnError = new AtomicBoolean(true);
AtomicBoolean isClientClosed = new AtomicBoolean(false);
AtomicInteger errorsBeforeClose = new AtomicInteger();
AtomicInteger errorsAfterClose = new AtomicInteger();
for (int i = 0; i < 500; ++i) {
WorkItemCommitRequest request = makeCommitRequest(i, i * (i < 480 ? 8 : 128));
commitRequestList.add(request);
commitRequests.put((long) i, request);
latches.add(new CountDownLatch(1));
}
Collections.shuffle(commitRequestList);
serviceRegistry.addService(
new CloudWindmillServiceV1Alpha1ImplBase() {
@Override
public StreamObserver<StreamingCommitWorkRequest> commitWorkStream(
StreamObserver<StreamingCommitResponse> responseObserver) {
StreamObserver<StreamingCommitWorkRequest> testCommitStreamObserver =
getTestCommitStreamObserver(responseObserver, commitRequests);
return new StreamObserver<StreamingCommitWorkRequest>() {
@Override
public void onNext(StreamingCommitWorkRequest request) {
if (shouldServerReturnError.get()) {
try {
responseObserver.onError(
new RuntimeException("shouldServerReturnError = true"));
if (isClientClosed.get()) {
errorsAfterClose.incrementAndGet();
} else {
errorsBeforeClose.incrementAndGet();
}
} catch (IllegalStateException e) {
}
} else {
testCommitStreamObserver.onNext(request);
}
}
@Override
public void onError(Throwable throwable) {
testCommitStreamObserver.onError(throwable);
}
@Override
public void onCompleted() {
testCommitStreamObserver.onCompleted();
}
};
}
});
CommitWorkStream stream = client.commitWorkStream();
for (int i = 0; i < commitRequestList.size(); ) {
final CountDownLatch latch = latches.get(i);
if (stream.commitWorkItem(
"computation",
commitRequestList.get(i),
(CommitStatus status) -> {
assertEquals(status, CommitStatus.OK);
latch.countDown();
})) {
i++;
} else {
stream.flush();
}
}
stream.flush();
long deadline = System.currentTimeMillis() + 60_000;
while (true) {
Thread.sleep(100);
int tmpErrorsBeforeClose = errorsBeforeClose.get();
if (tmpErrorsBeforeClose > 0) {
break;
}
if (System.currentTimeMillis() > deadline) {
fail(
String.format(
"Expected errors not sent by server errorsBeforeClose: %s"
+ " \n Should not reach here if the test is working as expected.",
tmpErrorsBeforeClose));
}
}
stream.close();
isClientClosed.set(true);
deadline = System.currentTimeMillis() + 60_000;
while (true) {
Thread.sleep(100);
int tmpErrorsAfterClose = errorsAfterClose.get();
if (tmpErrorsAfterClose > 0) {
break;
}
if (System.currentTimeMillis() > deadline) {
fail(
String.format(
"Expected errors not sent by server errorsAfterClose: %s"
+ " \n Should not reach here if the test is working as expected.",
tmpErrorsAfterClose));
}
}
shouldServerReturnError.set(false);
for (CountDownLatch latch : latches) {
assertTrue(latch.await(1, TimeUnit.MINUTES));
}
assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
}