test + fixes

Change-Id: Id4c56656a829f5f4c7ab1170f5f980cf3cc3760c

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java

@@ -47,6 +47,7 @@
 class BigtableUnaryOperationCallable<ReqT, RespT> extends UnaryCallable<ReqT, RespT> {
   private static final Logger LOGGER =
       Logger.getLogger(BigtableUnaryOperationCallable.class.getName());
+  Logger logger = LOGGER;
 
   private final ServerStreamingCallable<ReqT, RespT> inner;
   private final ApiCallContext defaultCallContext;
@@ -125,8 +126,8 @@ public void onResponse(RespT resp) {
         String msg =
             String.format(
                 "Received multiple responses for a %s unary operation. Previous: %s, New: %s",
-                spanName, resp, response);
-        LOGGER.log(Level.WARNING, msg);
+                spanName, response, resp);
+        logger.log(Level.WARNING, msg);
 
         InternalException error =
             new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
@@ -140,7 +141,6 @@ public void onResponse(RespT resp) {
 
       responseReceived = true;
       this.response = resp;
-      this.tracer.responseReceived();
     }
 
     @Override

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -1408,6 +1408,8 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
     ServerStreamingCallable<ReqT, RespT> transformed =
         new TransformingServerStreamingCallable<>(base, requestTransformer, responseTranformer);
 
+    transformed = new BigtableTracerStreamingCallable<>(transformed);
+
     return new BigtableUnaryOperationCallable<>(
         transformed,
         clientContext.getDefaultCallContext(),

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java

@@ -0,0 +1,149 @@
+package com.google.cloud.bigtable.data.v2.stub;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.InternalException;
+import com.google.api.gax.tracing.ApiTracerFactory;
+import com.google.api.gax.tracing.SpanName;
+import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable;
+import com.google.common.collect.ImmutableList;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.util.concurrent.ExecutionException;
+import java.util.logging.Logger;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+@RunWith(JUnit4.class)
+public class BigtableUnaryOperationCallableTest {
+    @Rule
+    public final MockitoRule mockitoRule = MockitoJUnit.rule();
+
+    @Mock
+    private ApiTracerFactory tracerFactory;
+    @Mock
+    private BigtableTracer tracer;
+
+    @Before
+    public void setUp() throws Exception {
+        Mockito.when(tracerFactory.newTracer(Mockito.any(), Mockito.any(), Mockito.any()))
+                .thenReturn(tracer);
+    }
+
+    @Test
+    public void testFutureResolve() throws Exception {
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                new FakeStreamingApi.ServerStreamingStashCallable<>(ImmutableList.of("value")),
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+
+        ApiFuture<String> f = callable.futureCall("fake");
+        assertThat(f.get()).isEqualTo("value");
+    }
+
+    @Test
+    public void testMultipleResponses() throws Exception {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        callable.logger = Mockito.mock(Logger.class);
+
+        ApiFuture<String> f = callable.futureCall("fake");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onResponse("first");
+        call.getController().getObserver().onResponse("second");
+
+        Throwable e = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+        assertThat(e).isInstanceOf(InternalException.class);
+        assertThat(e).hasMessageThat().contains("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+        ArgumentCaptor<String> msgCaptor = ArgumentCaptor.forClass(String.class);
+        verify(callable.logger).log(Mockito.any(), msgCaptor.capture());
+        assertThat(msgCaptor.getValue()).isEqualTo("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+        assertThat(call.getController().isCancelled()).isTrue();
+
+    }
+
+    @Test
+    public void testCancel() {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        f.cancel(true);
+
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        assertThat(call.getController().isCancelled()).isTrue();
+    }
+
+    @Test
+    public void testMissingResponse() {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onComplete();
+
+        Throwable cause = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+        assertThat(cause).hasMessageThat().isEqualTo("Fake.method unary operation completed without a response message");
+    }
+
+    @Test
+    public void testTracing() throws Exception {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onResponse("value");
+        call.getController().getObserver().onComplete();
+
+        f.get();
+        verify(tracer).responseReceived();
+        verify(tracer).operationSucceeded();
+
+        // afterResponse is the responsibility of BigtableTracerStreamingCallable
+        verify(tracer, never()).afterResponse(Mockito.anyLong());
+    }
+}