Add tests for call.cancel() from within messageRead.

Author: buchgr

core/src/test/java/io/grpc/internal/ClientCallImplTest.java

@@ -541,6 +541,39 @@ public void timeoutShouldNotBeSet() {
     assertFalse(headers.containsKey(GrpcUtil.TIMEOUT_KEY));
   }
 
+  @Test
+  public void cancelInOnMessageShouldInvokeStreamCancel() throws Exception {
+    final ClientCallImpl<Void, Void> call = new ClientCallImpl<Void, Void>(
+        DESCRIPTOR,
+        MoreExecutors.directExecutor(),
+        CallOptions.DEFAULT,
+        provider,
+        deadlineCancellationExecutor);
+    final Exception cause = new Exception();
+    ClientCall.Listener<Void> callListener =
+        new ClientCall.Listener<Void>() {
+          @Override
+          public void onMessage(Void message) {
+            call.cancel("foo", cause);
+          }
+        };
+
+    call.start(callListener, new Metadata());
+    call.halfClose();
+    call.request(1);
+
+    verify(stream).start(listenerArgumentCaptor.capture());
+    ClientStreamListener streamListener = listenerArgumentCaptor.getValue();
+    streamListener.onReady();
+    streamListener.headersRead(new Metadata());
+    streamListener.messageRead(new ByteArrayInputStream(new byte[0]));
+    verify(stream).cancel(statusCaptor.capture());
+    Status status = statusCaptor.getValue();
+    assertEquals(Status.CANCELLED.getCode(), status.getCode());
+    assertEquals("foo", status.getDescription());
+    assertSame(cause, status.getCause());
+  }
+
   private static class TestMarshaller<T> implements Marshaller<T> {
     @Override
     public InputStream stream(T value) {

testing/src/main/java/io/grpc/internal/testing/AbstractTransportTest.java

@@ -587,6 +587,59 @@ public void clientCancel() throws Exception {
     verify(mockClientStreamListener, never()).closed(any(Status.class), any(Metadata.class));
   }
 
+  @Test(timeout = 5000)
+  public void clientCancelFromWithinMessageRead() throws Exception {
+    server.start(serverListener);
+    client.start(mockClientTransportListener);
+    MockServerTransportListener serverTransportListener
+        = serverListener.takeListenerOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
+    serverTransport = serverTransportListener.transport;
+
+    final SettableFuture<Boolean> closedCalled = SettableFuture.create();
+    final ClientStream clientStream = client.newStream(methodDescriptor, new Metadata());
+    clientStream.start(new ClientStreamListener() {
+      @Override
+      public void headersRead(Metadata headers) {
+      }
+
+      @Override
+      public void closed(Status status, Metadata trailers) {
+        assertEquals(Status.CANCELLED.getCode(), status.getCode());
+        assertEquals("nevermind", status.getDescription());
+        closedCalled.set(true);
+      }
+
+      @Override
+      public void messageRead(InputStream message) {
+        assertEquals("foo", methodDescriptor.parseResponse(message));
+        clientStream.cancel(Status.CANCELLED.withDescription("nevermind"));
+      }
+
+      @Override
+      public void onReady() {
+      }
+    });
+    clientStream.halfClose();
+    clientStream.request(1);
+
+    StreamCreation serverStreamCreation
+        = serverTransportListener.takeStreamOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
+    assertEquals(methodDescriptor.getFullMethodName(), serverStreamCreation.method);
+    ServerStream serverStream = serverStreamCreation.stream;
+    ServerStreamListener mockServerStreamListener = serverStreamCreation.listener;
+    verify(mockServerStreamListener, timeout(TIMEOUT_MS)).onReady();
+
+    assertTrue(serverStream.isReady());
+    serverStream.writeHeaders(new Metadata());
+    serverStream.writeMessage(methodDescriptor.streamRequest("foo"));
+    serverStream.flush();
+
+    // Block until closedCalled was set.
+    closedCalled.get();
+
+    serverStream.close(Status.OK, new Metadata());
+  }
+
   @Test
   public void serverCancel() throws Exception {
     server.start(serverListener);