TracedBatchingCallable extends UnaryCallable<RequestT, ResponseT> { @Override public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) { ApiTracer tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.Batching); TraceFinisher<ResponseT> finisher = new TraceFinisher<>(tracer); try { long elementCount = batchingDescriptor.countElements(request); long requestSize = batchingDescriptor.countBytes(request); tracer.batchRequestSent(elementCount, requestSize); context = context.withTracer(tracer); ApiFuture<ResponseT> future = innerCallable.futureCall(request, context); ApiFutures.addCallback(future, finisher, MoreExecutors.directExecutor()); return future; } catch (RuntimeException e) { finisher.onFailure(e); throw e; } } TracedBatchingCallable(
      UnaryCallable<RequestT, ResponseT> innerCallable,
      ApiTracerFactory tracerFactory,
      SpanName spanName,
      BatchingDescriptor<RequestT, ResponseT> batchingDescriptor); @Override ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context);  }
@Test public void testRootTracerCreated() { tracedBatchingCallable.futureCall("test", callContext); verify(tracerFactory, times(1)) .newTracer(callContext.getTracer(), SPAN_NAME, OperationType.Batching); }
@Test public void testBatchAttributesStamped() { when(batchingDescriptor.countElements(anyString())).thenReturn(10L); when(batchingDescriptor.countBytes(anyString())).thenReturn(20L); tracedBatchingCallable.futureCall("test", callContext); verify(tracer).batchRequestSent(10, 20); }
@Test public void testOperationFinish() { innerResult.set("successful result"); tracedBatchingCallable.futureCall("test", callContext); verify(tracer, times(1)).operationSucceeded(); }
@Test public void testOperationFailed() { RuntimeException fakeError = new RuntimeException("fake error"); innerResult.setException(fakeError); tracedBatchingCallable.futureCall("test", callContext); verify(tracer, times(1)).operationFailed(fakeError); }
@Test public void testSyncError() { RuntimeException fakeError = new RuntimeException("fake error"); @SuppressWarnings("unchecked") UnaryCallable<String, String>[] innerCallableWrapper = new UnaryCallable[] {innerCallable}; reset(innerCallableWrapper); when(innerCallable.futureCall(eq("failing test"), any(ApiCallContext.class))) .thenThrow(fakeError); try { tracedBatchingCallable.futureCall("failing test", callContext); } catch (RuntimeException e) { } verify(tracer, times(1)).operationFailed(fakeError); }