diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/TracerProvider.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/TracerProvider.java index f39076bf10e7..526d657154fe 100644 --- a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/TracerProvider.java +++ b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/TracerProvider.java @@ -108,6 +108,20 @@ public Context extractContext(String diagnosticId, Context context) { return local; } + /** + * For each tracer plugged into the SDK a new context containing the span builder is returned. + * + * @param context Additional metadata containing the span name for creating the span builer. + */ + public Context getSharedSpanBuilder(Context context) { + Context local = Objects.requireNonNull(context, "'context' cannot be null."); + String spanName = getSpanName(ProcessKind.SEND); + for (Tracer tracer : tracers) { + local = tracer.getSharedSpanBuilder(spanName, local); + } + return local; + } + private void end(String statusMessage, Throwable throwable, Context context) { for (Tracer tracer : tracers) { tracer.end(statusMessage, throwable, context); diff --git a/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/TracerProviderTest.java b/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/TracerProviderTest.java index 1c48c1c00181..e9d3f120e6d0 100644 --- a/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/TracerProviderTest.java +++ b/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/TracerProviderTest.java @@ -24,6 +24,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.never; @@ -194,4 +195,40 @@ public void addSpanLinks() { // Act assertThrows(NullPointerException.class, () -> tracerProvider.addSpanLinks(null)); } + + @Test + public void getSpanBuilderReturnsUpdatedContext() { + // Arrange + final String spanBuilderKey = "spanBuilder-key"; + final String spanBuilderValue = "spanBuilder-value"; + + final String spanBuilderKey1 = "spanBuilder-key1"; + final String spanBuilderValue1 = "spanBuilder-value1"; + final Context startingContext = Context.NONE; + + when(tracer.getSharedSpanBuilder(anyString(), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(spanBuilderKey, spanBuilderValue); + } + ); + when(tracer2.getSharedSpanBuilder(anyString(), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(spanBuilderKey1, spanBuilderValue1); + } + ); + + // Act + final Context updatedContext = tracerProvider.getSharedSpanBuilder(startingContext); + + // Assert + final Optional spanBuilderData = updatedContext.getData(spanBuilderKey); + Assertions.assertTrue(spanBuilderData.isPresent()); + Assertions.assertEquals(spanBuilderValue, spanBuilderData.get()); + + final Optional spanBuilderData1 = updatedContext.getData(spanBuilderKey1); + Assertions.assertTrue(spanBuilderData1.isPresent()); + Assertions.assertEquals(spanBuilderValue1, spanBuilderData1.get()); + } } diff --git a/sdk/core/azure-core-tracing-opencensus/src/main/java/com/azure/core/tracing/opencensus/OpenCensusTracer.java b/sdk/core/azure-core-tracing-opencensus/src/main/java/com/azure/core/tracing/opencensus/OpenCensusTracer.java index c786fdfdbfe7..87e26d840623 100644 --- a/sdk/core/azure-core-tracing-opencensus/src/main/java/com/azure/core/tracing/opencensus/OpenCensusTracer.java +++ b/sdk/core/azure-core-tracing-opencensus/src/main/java/com/azure/core/tracing/opencensus/OpenCensusTracer.java @@ -181,6 +181,13 @@ public Context extractContext(String diagnosticId, Context context) { return AmqpPropagationFormatUtil.extractContext(diagnosticId, context); } + @Override + public Context getSharedSpanBuilder(String spanName, Context context) { + throw logger.logExceptionAsError( + new UnsupportedOperationException("This operation is not supported in OpenCensus")); + // Remove OpenCensus support for tracing in https://github.com/Azure/azure-sdk-for-java/issues/6781 + } + /** * Starts a new child {@link Span} with parent being the remote and uses the {@link Span} is in the current Context, * to return an object that represents that scope. diff --git a/sdk/core/azure-core-tracing-opentelemetry/CHANGELOG.md b/sdk/core/azure-core-tracing-opentelemetry/CHANGELOG.md index 780d77ac2ad6..3fd77148407c 100644 --- a/sdk/core/azure-core-tracing-opentelemetry/CHANGELOG.md +++ b/sdk/core/azure-core-tracing-opentelemetry/CHANGELOG.md @@ -1,7 +1,8 @@ # Release History ## Version 1.0.0-beta.2 (2020-01-03) -- Add eventhub properties to attributes of processing spans. +- Add `EventHubs.*` properties to attributes of processing spans. - Remove `Azure` prefix from convenience layer span names. +- Add links for batch send operation in Event Hubs client library. This package's [documentation](https://github.com/Azure/azure-sdk-for-java/blob/azure-core-tracing-opentelemetry_1.0.0-beta.2/sdk/core/azure-core-tracing-opentelemetry/README.md) @@ -19,4 +20,3 @@ This package's [documentation](https://github.com/Azure/azure-sdk-for-java/blob/azure-core-tracing-opentelemetry_1.0.0-beta.1/sdk/core/azure-core-tracing-opentelemetry/README.md) and [samples](https://github.com/Azure/azure-sdk-for-java/blob/azure-core-tracing-opentelemetry_1.0.0-beta.1/sdk/core/azure-core-tracing-opentelemetry/src/samples). - diff --git a/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java b/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java index d0040fcf23fc..51cfe65b84ee 100644 --- a/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java +++ b/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java @@ -66,7 +66,11 @@ public Context start(String spanName, Context context, ProcessKind processKind) switch (processKind) { case SEND: - spanBuilder = getSpanBuilder(spanName, context); + // use previously created span builder from the LINK process. + spanBuilder = getOrDefault(context, SPAN_BUILDER_KEY, null, Builder.class); + if (spanBuilder == null) { + return Context.NONE; + } span = spanBuilder.setSpanKind(Span.Kind.PRODUCER).startSpan(); if (span.isRecording()) { // If span is sampled in, add additional request attributes @@ -172,6 +176,11 @@ public Context extractContext(String diagnosticId, Context context) { return AmqpPropagationFormatUtil.extractContext(diagnosticId, context); } + @Override + public Context getSharedSpanBuilder(String spanName, Context context) { + return context.addData(SPAN_BUILDER_KEY, getSpanBuilder(spanName, context)); + } + /** * Starts a new child {@link Span} with parent being the remote and uses the {@link Span} is in the current Context, * to return an object that represents that scope. diff --git a/sdk/core/azure-core-tracing-opentelemetry/src/samples/PublishEventsTracingSample.md b/sdk/core/azure-core-tracing-opentelemetry/src/samples/PublishEventsTracingSample.md index 2bd9d1bd6ab7..6256f61fba68 100644 --- a/sdk/core/azure-core-tracing-opentelemetry/src/samples/PublishEventsTracingSample.md +++ b/sdk/core/azure-core-tracing-opentelemetry/src/samples/PublishEventsTracingSample.md @@ -60,6 +60,7 @@ public class Sample { private static final Logger LOGGER = getLogger("Sample"); private static final Tracer TRACER; private static final TracerSdkFactory TRACER_SDK_FACTORY; + private static final Duration OPERATION_TIMEOUT = Duration.ofSeconds(30); static { TRACER_SDK_FACTORY = configureOpenTelemetryAndLoggingExporter(); @@ -80,35 +81,66 @@ public class Sample { } private static void doClientWork() { - EventHubProducerClient producer = new EventHubClientBuilder() - .connectionString(connectionString) - .buildProducerClient(); + EventHubProducerAsyncClient producer = new EventHubClientBuilder() + .connectionString(CONNECTION_STRING) + .buildAsyncProducerClient(); Span span = TRACER.spanBuilder("user-parent-span").startSpan(); try (final Scope scope = TRACER.withSpan(span)) { - final EventData event1 = new EventData("1".getBytes(UTF_8)); - event1.addContext(PARENT_SPAN_KEY, span); - - final EventData event2 = new EventData("2".getBytes(UTF_8)); - event2.addContext(PARENT_SPAN_KEY, span); + String firstPartition = producer.getPartitionIds().blockFirst(OPERATION_TIMEOUT); + + final byte[] body = "EventData Sample 1".getBytes(UTF_8); + final byte[] body2 = "EventData Sample 2".getBytes(UTF_8); - final List telemetryEvents = Arrays.asList(event1, event2); + // We will publish three events based on simple sentences. + Flux data = Flux.just( + new EventData(body).addContext(PARENT_SPAN_KEY, TRACER.getCurrentSpan()), + new EventData(body2).addContext(PARENT_SPAN_KEY, TRACER.getCurrentSpan())); + + // Create a batch to send the events. final CreateBatchOptions options = new CreateBatchOptions() - .setPartitionKey("telemetry") + .setPartitionId(firstPartition) .setMaximumSizeInBytes(256); - - EventDataBatch currentBatch = producer.createBatch(options); - - // For each telemetry event, we try to add it to the current batch. - for (EventData event : telemetryEvents) { - if (!currentBatch.tryAdd(event)) { - producer.send(currentBatch); - currentBatch = producer.createBatch(options); + + final AtomicReference currentBatch = new AtomicReference<>( + producer.createBatch(options).block(OPERATION_TIMEOUT)); + + data.flatMap(event -> { + final EventDataBatch batch = currentBatch.get(); + if (batch.tryAdd(event)) { + return Mono.empty(); } - } - } finally { - span.end(); - producer.close(); + + // The batch is full, so we create a new batch and send the batch. Mono.when completes when both operations + // have completed. + return Mono.when( + producer.send(batch), + producer.createBatch(options).map(newBatch -> { + currentBatch.set(newBatch); + + // Add that event that we couldn't before. + if (!newBatch.tryAdd(event)) { + throw Exceptions.propagate(new IllegalArgumentException(String.format( + "Event is too large for an empty batch. Max size: %s. Event: %s", + newBatch.getMaxSizeInBytes(), event.getBodyAsString()))); + } + + return newBatch; + })); + }).then() + .doFinally(signal -> { + final EventDataBatch batch = currentBatch.getAndSet(null); + if (batch != null) { + producer.send(batch).block(OPERATION_TIMEOUT); + } + }) + .subscribe(unused -> System.out.println("Complete"), + error -> System.out.println("Error sending events: " + error), + () -> { + span.end(); + producer.close(); + System.out.println("Completed sending events."); + }); } } } diff --git a/sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java b/sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java index dfb362de9756..0103a549707a 100644 --- a/sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java +++ b/sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java @@ -92,7 +92,6 @@ public void startSpanParentContextFlowTest() { final ReadableSpan recordEventsSpan = (ReadableSpan) updatedContext.getData(PARENT_SPAN_KEY).get(); assertEquals(Span.Kind.INTERNAL, recordEventsSpan.toSpanData().getKind()); - assertEquals(true, true); } @Test @@ -117,9 +116,11 @@ public void startSpanTestNoUserParent() { public void startSpanProcessKindSend() { // Arrange final SpanId parentSpanId = parentSpan.getContext().getSpanId(); + // Start user parent span. + final Span.Builder spanBuilder = tracer.spanBuilder(METHOD_NAME); // Add additional metadata to spans for SEND final Context traceContext = tracingContext.addData(ENTITY_PATH_KEY, ENTITY_PATH_VALUE) - .addData(HOST_NAME_KEY, HOSTNAME_VALUE); + .addData(HOST_NAME_KEY, HOSTNAME_VALUE).addData(SPAN_BUILDER_KEY, spanBuilder); // Act final Context updatedContext = openTelemetryTracer.start(METHOD_NAME, traceContext, ProcessKind.SEND); @@ -181,6 +182,14 @@ public void startSpanProcessKindProcess() { verifySpanAttributes(attributeMap); } + @Test + public void getSpanBuilderTest() { + // Act + final Context updatedContext = openTelemetryTracer.getSharedSpanBuilder(METHOD_NAME, Context.NONE); + + assertTrue(updatedContext.getData(SPAN_BUILDER_KEY).isPresent()); + } + @Test public void startProcessSpanWithRemoteParent() { // Arrange @@ -277,7 +286,7 @@ public void endSpanTestThrowableResponseCode() { } private static void assertSpanWithExplicitParent(Context updatedContext, SpanId parentSpanId) { - assertNotNull(updatedContext.getData(PARENT_SPAN_KEY)); + assertNotNull(updatedContext.getData(PARENT_SPAN_KEY).get()); // verify instance created of opentelemetry-sdk (test impl), span implementation assertTrue(updatedContext.getData(PARENT_SPAN_KEY).get() instanceof ReadableSpan); @@ -292,7 +301,7 @@ private static void assertSpanWithExplicitParent(Context updatedContext, SpanId } private static void assertSpanWithRemoteParent(Context updatedContext, SpanId parentSpanId) { - assertNotNull(updatedContext.getData(PARENT_SPAN_KEY)); + assertNotNull(updatedContext.getData(PARENT_SPAN_KEY).get()); // verify instance created of opentelemetry-sdk (test impl), span implementation assertTrue(updatedContext.getData(PARENT_SPAN_KEY).get() instanceof ReadableSpan); diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/util/tracing/Tracer.java b/sdk/core/azure-core/src/main/java/com/azure/core/util/tracing/Tracer.java index 463450338325..acf30ddcc7cb 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/util/tracing/Tracer.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/util/tracing/Tracer.java @@ -204,4 +204,23 @@ public interface Tracer { * @throws NullPointerException if {@code diagnosticId} or {@code context} is {@code null}. */ Context extractContext(String diagnosticId, Context context); + + /** + * Returns a span builder with the provided name in {@link Context}. + * + *

Code samples

+ * + *

Returns a builder with the provided span name.

+ * {@codesnippet com.azure.core.util.tracing.getSharedSpanBuilder#string-context} + * + * @param spanName Name to give the span for the created builder. + * @param context Additional metadata that is passed through the call stack. + * + * @return The updated {@link Context} object containing the span builder. + * @throws NullPointerException if {@code context} or {@code spanName} is {@code null}. + */ + default Context getSharedSpanBuilder(String spanName, Context context) { + // no-op + return Context.NONE; + } } diff --git a/sdk/core/azure-core/src/samples/java/com/azure/core/util/tracing/TracerJavaDocCodeSnippets.java b/sdk/core/azure-core/src/samples/java/com/azure/core/util/tracing/TracerJavaDocCodeSnippets.java index b84e833acbd8..5ae137b2be54 100644 --- a/sdk/core/azure-core/src/samples/java/com/azure/core/util/tracing/TracerJavaDocCodeSnippets.java +++ b/sdk/core/azure-core/src/samples/java/com/azure/core/util/tracing/TracerJavaDocCodeSnippets.java @@ -7,6 +7,7 @@ import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_BUILDER_KEY; /** * Contains code snippets when generating javadocs through doclets for {@link Tracer}. @@ -120,6 +121,18 @@ public void extractContext() { // END: com.azure.core.util.tracing.extractContext#string-context } + /** + * Code snippet for {@link Tracer#getSharedSpanBuilder(String, Context)} + */ + public void getSharedSpanBuilder() { + // BEGIN: com.azure.core.util.tracing.getSpanBuilder#string-context + // Returns a span builder with the provided name + String methodName = "message-span"; + Context spanContext = tracer.getSharedSpanBuilder(methodName, Context.NONE); + System.out.printf("Span context of the current tracing span: %s%n", spanContext.getData(SPAN_BUILDER_KEY).get()); + // END: com.azure.core.util.tracing.getSpanBuilder#string-context + } + //Noop Tracer private static final class TracerImplementation implements Tracer { @Override @@ -161,5 +174,10 @@ public void addLink(Context context) { public Context extractContext(String diagnosticId, Context context) { return null; } + + @Override + public Context getSharedSpanBuilder(String spanName, Context context) { + return null; + } } } diff --git a/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java b/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java index 3483c51ae712..6d394d2c414d 100644 --- a/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java +++ b/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java @@ -254,7 +254,7 @@ Context getContext() { * @throws NullPointerException if {@code key} or {@code value} is null. * @return The updated {@link EventData}. */ - EventData addContext(String key, Object value) { + public EventData addContext(String key, Object value) { Objects.requireNonNull(key, "The 'key' parameter cannot be null."); Objects.requireNonNull(value, "The 'value' parameter cannot be null."); this.context = context.addData(key, value); diff --git a/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClient.java b/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClient.java index 1b68062096f2..2c8f37fb2a4a 100644 --- a/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClient.java +++ b/sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClient.java @@ -42,11 +42,11 @@ import java.util.function.Function; import java.util.function.Supplier; import java.util.stream.Collector; -import java.util.stream.Collectors; import static com.azure.core.util.FluxUtil.monoError; import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static com.azure.messaging.eventhubs.implementation.ClientConstants.MAX_MESSAGE_LENGTH_BYTES; /** @@ -397,11 +397,20 @@ public Mono send(EventDataBatch batch) { ? new AtomicReference<>(Context.NONE) : null; - final List messages = batch.getEvents().stream().map(event -> { - final Message message = messageSerializer.serialize(event); + Context sharedContext = null; + List messages = new ArrayList<>(); + + for (int i = 0; i < batch.getEvents().size(); i++) { + final EventData event = batch.getEvents().get(i); if (isTracingEnabled) { parentContext.set(event.getContext()); + if (i == 0) { + sharedContext = tracerProvider.getSharedSpanBuilder(parentContext.get()); + } + tracerProvider.addSpanLinks(sharedContext.addData(SPAN_CONTEXT_KEY, event.getContext())); } + final Message message = messageSerializer.serialize(event); + if (!CoreUtils.isNullOrEmpty(partitionKey)) { final MessageAnnotations messageAnnotations = message.getMessageAnnotations() == null ? new MessageAnnotations(new HashMap<>()) @@ -409,15 +418,14 @@ public Mono send(EventDataBatch batch) { messageAnnotations.getValue().put(AmqpConstants.PARTITION_KEY, partitionKey); message.setMessageAnnotations(messageAnnotations); } + messages.add(message); + } - return message; - }).collect(Collectors.toList()); - + Context finalSharedContext = sharedContext; return getSendLink(batch.getPartitionId()) .flatMap(link -> { if (isTracingEnabled) { - Context userSpanContext = parentContext.get(); - Context entityContext = userSpanContext.addData(ENTITY_PATH_KEY, link.getEntityPath()); + Context entityContext = finalSharedContext.addData(ENTITY_PATH_KEY, link.getEntityPath()); // start send span and store updated context parentContext.set(tracerProvider.startSpan( entityContext.addData(HOST_NAME_KEY, link.getHostname()), ProcessKind.SEND)); diff --git a/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClientTest.java b/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClientTest.java index 0676d5f1145f..198cb77c8407 100644 --- a/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClientTest.java +++ b/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerAsyncClientTest.java @@ -44,6 +44,7 @@ import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_BUILDER_KEY; import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static java.nio.charset.StandardCharsets.UTF_8; import static org.mockito.ArgumentMatchers.any; @@ -244,6 +245,13 @@ public void sendStartSpanSingleMessage() { } ); + when(tracer1.getSharedSpanBuilder(eq("EventHubs.send"), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(SPAN_BUILDER_KEY, "value"); + } + ); + // Act StepVerifier.create(asyncProducer.send(testData, sendOptions)) .verifyComplete(); @@ -294,6 +302,13 @@ public void sendMessageRetrySpanTest() { } ); + when(tracer1.getSharedSpanBuilder(eq("EventHubs.send"), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(SPAN_BUILDER_KEY, "value"); + } + ); + //Act StepVerifier.create(asyncProducer.send(testData, sendOptions)).verifyComplete(); @@ -301,7 +316,7 @@ public void sendMessageRetrySpanTest() { verify(tracer1, times(1)) .start(eq("EventHubs.send"), any(), eq(ProcessKind.SEND)); verify(tracer1, never()).start(eq("EventHubs.message"), any(), eq(ProcessKind.MESSAGE)); - verify(tracer1, never()).addLink(any()); + verify(tracer1, times(2)).addLink(any()); verify(tracer1, times(1)).end(eq("success"), isNull(), any()); } diff --git a/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerClientTest.java b/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerClientTest.java index b06267314e79..246753b9273c 100644 --- a/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerClientTest.java +++ b/sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventHubProducerClientTest.java @@ -44,6 +44,7 @@ import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_BUILDER_KEY; import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static java.nio.charset.StandardCharsets.UTF_8; import static org.mockito.ArgumentMatchers.any; @@ -168,6 +169,13 @@ public void sendStartSpanSingleMessage() { return passed.addData(PARENT_SPAN_KEY, "value").addData(DIAGNOSTIC_ID_KEY, "value2"); } ); + when(tracer1.getSharedSpanBuilder(eq("EventHubs.send"), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(SPAN_BUILDER_KEY, "value"); + } + ); + //Act producer.send(eventData); @@ -209,13 +217,20 @@ public void sendMessageRetrySpanTest() { } ); + when(tracer1.getSharedSpanBuilder(eq("EventHubs.send"), any())).thenAnswer( + invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(SPAN_BUILDER_KEY, "value"); + } + ); + //Act producer.send(eventData); //Assert verify(tracer1, times(1)).start(eq("EventHubs.send"), any(), eq(ProcessKind.SEND)); verify(tracer1, never()).start(eq("EventHubs.message"), any(), eq(ProcessKind.MESSAGE)); - verify(tracer1, never()).addLink(any()); + verify(tracer1, times(1)).addLink(any()); verify(tracer1, times(1)).end(eq("success"), isNull(), any()); }