fix: set OperationType on traced callables context

Author: diegomarquezp

gax-java/gax/src/main/java/com/google/api/gax/tracing/MetricsTracerFactory.java

@@ -31,6 +31,7 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 
@@ -66,6 +67,7 @@ public MetricsTracerFactory(MetricsRecorder metricsRecorder, Map<String, String>
 
   @Override
   public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType operationType) {
+    Preconditions.checkNotNull(operationType, "operationType cannot be null.");
     MetricsTracer metricsTracer =
         new MetricsTracer(
             MethodName.of(spanName.getClientName(), spanName.getMethodName()), metricsRecorder);

gax-java/gax/src/main/java/com/google/api/gax/tracing/OpencensusTracerFactory.java

@@ -96,6 +96,7 @@ public OpencensusTracerFactory(Map<String, String> spanAttributes) {
   /** {@inheritDoc } */
   @Override
   public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType operationType) {
+    Preconditions.checkNotNull(operationType, "operationType cannot be null.");
     // Default to the current in context span. This is used for outermost tracers that inherit
     // the caller's parent span.
     Span parentSpan = internalTracer.getCurrentSpan();

gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java

@@ -73,7 +73,8 @@ public TracedBatchingCallable(
       ApiTracerContext apiTracerContext,
       BatchingDescriptor<RequestT, ResponseT> batchingDescriptor) {
     this.tracerFactory = tracerFactory;
-    this.apiTracerContext = apiTracerContext;
+    this.apiTracerContext =
+        apiTracerContext.toBuilder().setOperationType(OperationType.Batching).build();
     this.batchingDescriptor = batchingDescriptor;
     this.innerCallable = innerCallable;
     this.spanName = SpanName.of(apiTracerContext);

gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBidiCallable.java

@@ -74,7 +74,9 @@ public TracedBidiCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.BidiStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }

gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedClientStreamingCallable.java

@@ -74,7 +74,9 @@ public TracedClientStreamingCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.ClientStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }

gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedServerStreamingCallable.java

@@ -70,7 +70,9 @@ public TracedServerStreamingCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.ClientStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }

gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java

@@ -68,7 +68,8 @@ public TracedUnaryCallable(
       ApiTracerContext apiTracerContext) {
     this.innerCallable = innerCallable;
     this.tracerFactory = tracerFactory;
-    this.apiTracerContext = apiTracerContext;
+    this.apiTracerContext =
+        apiTracerContext.toBuilder().setOperationType(OperationType.Unary).build();
     this.spanName = SpanName.of(apiTracerContext);
   }
 

gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedServerStreamingCallableTest.java

@@ -40,7 +40,6 @@
 
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.CancelledException;
-import com.google.api.gax.rpc.LibraryMetadata;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode.Code;
@@ -49,7 +48,6 @@
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockResponseObserver;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCall;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCallable;
-import com.google.api.gax.tracing.ApiTracerContext.Transport;
 import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -60,16 +58,11 @@
 @ExtendWith(MockitoExtension.class)
 class TracedServerStreamingCallableTest {
   private static final SpanName SPAN_NAME = SpanName.of("FakeClient", "FakeRpc");
-  private static final ApiTracerContext TRACER_CONTEXT =
-      ApiTracerContext.newBuilder()
-          .setFullMethodName("FakeClient/FakeRpc")
-          .setTransport(Transport.GRPC)
-          .setLibraryMetadata(LibraryMetadata.empty())
-          .setOperationType(OperationType.ServerStreaming)
-          .build();
 
   @Mock private ApiTracerFactory tracerFactory;
   @Mock private ApiTracer tracer;
+  @Mock private ApiTracerContext tracerContext;
+  @Mock private ApiTracerContext.Builder tracerContextBuilder;
 
   private MockServerStreamingCallable<String, String> innerCallable;
   private TracedServerStreamingCallable<String, String> tracedCallable;
@@ -81,10 +74,14 @@ void init(boolean useContext) {
     innerCallable = new MockServerStreamingCallable<>();
     // Wire the mock tracer factory
     if (useContext) {
+      when(tracerContext.toBuilder()).thenReturn(tracerContextBuilder);
+      when(tracerContextBuilder.setOperationType(any())).thenReturn(tracerContextBuilder);
+      when(tracerContextBuilder.build()).thenReturn(tracerContext);
       when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
-      tracedCallable =
-          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
+      when(tracerContext.fullMethodName()).thenReturn("FakeClient/FakeRpc");
+      t statutracedCallable =
+          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, tracerContext);
     } else {
       when(tracerFactory.newTracer(
               any(ApiTracer.class), any(SpanName.class), eq(OperationType.ServerStreaming)))
@@ -102,7 +99,7 @@ void testTracerCreated(boolean useContext) {
     init(useContext);
     tracedCallable.call("test", responseObserver, callContext);
     if (useContext) {
-      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, tracerContext);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(parentTracer, SPAN_NAME, OperationType.ServerStreaming);
@@ -189,7 +186,7 @@ void testSyncError(boolean useContext) {
     // Recreate the tracedCallable using the new inner callable
     if (useContext) {
       tracedCallable =
-          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
+          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, tracerContext);
     } else {
       tracedCallable = new TracedServerStreamingCallable<>(innerCallable, tracerFactory, SPAN_NAME);
     }