From 1c20eb6cef82e89ca4b67ef9f15105b187ee239d Mon Sep 17 00:00:00 2001
From: ejona <ejona@google.com>
Date: Tue, 23 Dec 2014 11:30:37 -0800
Subject: [PATCH] Implement cancel_after_* integration tests

-------------
Created by MOE: http://code.google.com/p/moe-java
MOE_MIGRATED_REVID=82734780
---
 .../integration/AbstractTransportTest.java    | 41 +++++++++++++++++++
 .../integration/TestServiceClient.java        |  4 ++
 2 files changed, 45 insertions(+)

diff --git a/integration-testing/src/main/java/com/google/net/stubby/testing/integration/AbstractTransportTest.java b/integration-testing/src/main/java/com/google/net/stubby/testing/integration/AbstractTransportTest.java
index d4400b208f..d7949db0ff 100644
--- a/integration-testing/src/main/java/com/google/net/stubby/testing/integration/AbstractTransportTest.java
+++ b/integration-testing/src/main/java/com/google/net/stubby/testing/integration/AbstractTransportTest.java
@@ -72,6 +72,7 @@ import org.junit.Assert;
 import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -305,6 +306,46 @@ public abstract class AbstractTransportTest {
     verifyNoMoreInteractions(responseObserver);
   }
 
+  @Test
+  public void cancelAfterBegin() throws Exception {
+    StreamRecorder<StreamingInputCallResponse> responseObserver = StreamRecorder.create();
+    StreamObserver<StreamingInputCallRequest> requestObserver =
+        asyncStub.streamingInputCall(responseObserver);
+    requestObserver.onError(new RuntimeException());
+    responseObserver.awaitCompletion();
+    assertEquals(Arrays.<StreamingInputCallResponse>asList(), responseObserver.getValues());
+    assertEquals(Status.CANCELLED, Status.fromThrowable(responseObserver.getError()));
+  }
+
+  @Test
+  public void cancelAfterFirstResponse() throws Exception {
+    final StreamingOutputCallRequest request = StreamingOutputCallRequest.newBuilder()
+        .addResponseParameters(ResponseParameters.newBuilder()
+            .setSize(31415))
+        .setPayload(Payload.newBuilder()
+            .setBody(ByteString.copyFrom(new byte[27182])))
+        .build();
+    final StreamingOutputCallResponse goldenResponse = StreamingOutputCallResponse.newBuilder()
+        .setPayload(Payload.newBuilder()
+            .setType(PayloadType.COMPRESSABLE)
+            .setBody(ByteString.copyFrom(new byte[31415])))
+        .build();
+
+    @SuppressWarnings("unchecked")
+    StreamObserver<StreamingOutputCallResponse> responseObserver = mock(StreamObserver.class);
+    StreamObserver<StreamingOutputCallRequest> requestObserver
+        = asyncStub.fullDuplexCall(responseObserver);
+    requestObserver.onValue(request);
+    verify(responseObserver, timeout(1000)).onValue(goldenResponse);
+    verifyNoMoreInteractions(responseObserver);
+
+    requestObserver.onError(new RuntimeException());
+    ArgumentCaptor<Throwable> captor = ArgumentCaptor.forClass(Throwable.class);
+    verify(responseObserver, timeout(1000)).onError(captor.capture());
+    assertEquals(Status.CANCELLED, Status.fromThrowable(captor.getValue()));
+    verifyNoMoreInteractions(responseObserver);
+  }
+
   @Test
   public void fullDuplexCallShouldSucceed() throws Exception {
     // Build the request.
diff --git a/integration-testing/src/main/java/com/google/net/stubby/testing/integration/TestServiceClient.java b/integration-testing/src/main/java/com/google/net/stubby/testing/integration/TestServiceClient.java
index cd04800b3a..241997d365 100644
--- a/integration-testing/src/main/java/com/google/net/stubby/testing/integration/TestServiceClient.java
+++ b/integration-testing/src/main/java/com/google/net/stubby/testing/integration/TestServiceClient.java
@@ -223,6 +223,10 @@ public class TestServiceClient {
       tester.pingPong();
     } else if ("empty_stream".equals(testCase)) {
       tester.emptyStream();
+    } else if ("cancel_after_begin".equals(testCase)) {
+      tester.cancelAfterBegin();
+    } else if ("cancel_after_first_response".equals(testCase)) {
+      tester.cancelAfterFirstResponse();
     } else {
       throw new IllegalArgumentException("Unknown test case: " + testCase);
     }