diff --git a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UrlReplaceHandler.java b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UrlReplaceHandler.java index cf2d2bbaa..75d881b87 100644 --- a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UrlReplaceHandler.java +++ b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UrlReplaceHandler.java @@ -61,6 +61,10 @@ public Response intercept(@Nonnull Chain chain) throws IOException { } try { request = replaceRequestUrl(request, replaceOption.getReplacementPairs()); + if (span != null) { + request = request.newBuilder().tag(Span.class, span).build(); + } + return chain.proceed(request); } finally { if (scope != null) { scope.close(); @@ -69,7 +73,6 @@ public Response intercept(@Nonnull Chain chain) throws IOException { span.end(); } } - return chain.proceed(request); } /** diff --git a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UserAgentHandler.java b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UserAgentHandler.java index 8e6eb7500..ec2831815 100644 --- a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UserAgentHandler.java +++ b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/UserAgentHandler.java @@ -73,6 +73,10 @@ public Response intercept(final Chain chain) throws IOException { } else if (!currentValue.contains(valueToAppend)) { builder.header(USER_AGENT_HEADER_NAME, currentValue + " " + valueToAppend); } + if (span != null) { + builder.tag(Span.class, span); + } + return chain.proceed(builder.build()); } finally { if (scope != null) { scope.close(); @@ -81,6 +85,5 @@ public Response intercept(final Chain chain) throws IOException { span.end(); } } - return chain.proceed(builder.build()); } } diff --git a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UrlReplaceHandlerTest.java b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UrlReplaceHandlerTest.java index 0c04f014e..4a6e88210 100644 --- a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UrlReplaceHandlerTest.java +++ b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UrlReplaceHandlerTest.java @@ -2,16 +2,25 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.*; import com.microsoft.kiota.http.KiotaClientFactory; import com.microsoft.kiota.http.middleware.options.UrlReplaceHandlerOption; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.Tracer; + import okhttp3.Interceptor; import okhttp3.OkHttpClient; import okhttp3.Request; import okhttp3.Response; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; import java.io.IOException; import java.util.HashMap; @@ -96,4 +105,92 @@ void testItDoesNotReplaceTokensInThePath() throws IOException { "https://graph.microsoft.com/v1.0/me/calendar/events/AAMkADM4Njk4YzMzLTc3YTktNGQwZi05ZjRjLTJlNTUzYWMxZjAxNABGAAAAAABzQLabgoJzT6ewNe8DeydoBwBoFwNVrXtLQ5PleemFoZ87AAAAAAENAABoFwNVrXtLQ5PleemFoZ87AAA33T-eAAA%3D", response.request().url().toString()); } + + @Test + void testUrlReplaceHandler_propagatesSpan() throws IOException { + try (MockedStatic otelMock = mockStatic(GlobalOpenTelemetry.class)) { + final Tracer mockTracer = mock(Tracer.class); + final SpanBuilder mockSpanBuilder = mock(SpanBuilder.class); + final Span mockSpan = mock(Span.class); + final io.opentelemetry.context.Scope mockScope = + mock(io.opentelemetry.context.Scope.class); + + otelMock.when(() -> GlobalOpenTelemetry.getTracer(anyString())).thenReturn(mockTracer); + when(mockTracer.spanBuilder(anyString())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.setParent(any())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.startSpan()).thenReturn(mockSpan); + when(mockSpan.makeCurrent()).thenReturn(mockScope); + + final com.microsoft.kiota.http.ObservabilityOptions obsOptions = + new com.microsoft.kiota.http.ObservabilityOptions(); + final Interceptor.Chain mockChain = mock(Interceptor.Chain.class); + final Response mockResponse = mock(Response.class); + defaultReplacementPairs.put("/users/TokenToReplace", "/me"); + final UrlReplaceHandler handler = + new UrlReplaceHandler(new UrlReplaceHandlerOption(defaultReplacementPairs)); + final Request request = + new Request.Builder() + .url(DEFAULT_URL_WITH_TOKEN) + .tag(com.microsoft.kiota.http.ObservabilityOptions.class, obsOptions) + .build(); + when(mockChain.request()).thenReturn(request); + when(mockChain.proceed(any(Request.class))) + .thenAnswer( + invocation -> { + final Request req = invocation.getArgument(0); + assertSame(mockSpan, req.tag(Span.class)); + verify(mockSpan, never()).end(); + verify(mockScope, never()).close(); + return mockResponse; + }); + + final Response response = handler.intercept(mockChain); + assertNotNull(response); + verify(mockScope).close(); + verify(mockSpan).end(); + } + } + + @Test + void testUrlReplaceHandler_cleansUpSpanOnException() throws IOException { + try (MockedStatic otelMock = mockStatic(GlobalOpenTelemetry.class)) { + final Tracer mockTracer = mock(Tracer.class); + final SpanBuilder mockSpanBuilder = mock(SpanBuilder.class); + final Span mockSpan = mock(Span.class); + final io.opentelemetry.context.Scope mockScope = + mock(io.opentelemetry.context.Scope.class); + + otelMock.when(() -> GlobalOpenTelemetry.getTracer(anyString())).thenReturn(mockTracer); + when(mockTracer.spanBuilder(anyString())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.setParent(any())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.startSpan()).thenReturn(mockSpan); + when(mockSpan.makeCurrent()).thenReturn(mockScope); + + final com.microsoft.kiota.http.ObservabilityOptions obsOptions = + new com.microsoft.kiota.http.ObservabilityOptions(); + final Interceptor.Chain mockChain = mock(Interceptor.Chain.class); + defaultReplacementPairs.put("/users/TokenToReplace", "/me"); + final UrlReplaceHandler handler = + new UrlReplaceHandler(new UrlReplaceHandlerOption(defaultReplacementPairs)); + final Request request = + new Request.Builder() + .url(DEFAULT_URL_WITH_TOKEN) + .tag(com.microsoft.kiota.http.ObservabilityOptions.class, obsOptions) + .build(); + when(mockChain.request()).thenReturn(request); + when(mockChain.proceed(any(Request.class))) + .thenAnswer( + invocation -> { + final Request req = invocation.getArgument(0); + assertSame(mockSpan, req.tag(Span.class)); + verify(mockSpan, never()).end(); + verify(mockScope, never()).close(); + throw new IOException("network failure"); + }); + + assertThrows(IOException.class, () -> handler.intercept(mockChain)); + verify(mockScope).close(); + verify(mockSpan).end(); + } + } } diff --git a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UserAgentHandlerTest.java b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UserAgentHandlerTest.java index 1b8fe4498..61421675b 100644 --- a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UserAgentHandlerTest.java +++ b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/UserAgentHandlerTest.java @@ -3,10 +3,17 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.*; import com.microsoft.kiota.http.middleware.options.UserAgentHandlerOption; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.Tracer; + import okhttp3.Headers; import okhttp3.Interceptor.Chain; import okhttp3.Request; @@ -14,6 +21,7 @@ import okhttp3.ResponseBody; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; @@ -84,4 +92,86 @@ void doesNotAddTheProductWhenDisabled() throws IOException { assertNotNull(result); assertNull(result.header("User-Agent")); } + + @Test + void enclosesProceedAndPropagatesSpan() throws IOException { + try (MockedStatic otelMock = mockStatic(GlobalOpenTelemetry.class)) { + final Tracer mockTracer = mock(Tracer.class); + final SpanBuilder mockSpanBuilder = mock(SpanBuilder.class); + final Span mockSpan = mock(Span.class); + final io.opentelemetry.context.Scope mockScope = + mock(io.opentelemetry.context.Scope.class); + + otelMock.when(() -> GlobalOpenTelemetry.getTracer(anyString())).thenReturn(mockTracer); + when(mockTracer.spanBuilder(anyString())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.setParent(any())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.startSpan()).thenReturn(mockSpan); + when(mockSpan.makeCurrent()).thenReturn(mockScope); + + final com.microsoft.kiota.http.ObservabilityOptions obsOptions = + new com.microsoft.kiota.http.ObservabilityOptions(); + final UserAgentHandler handler = new UserAgentHandler(); + final Request request = + new Request.Builder() + .url("http://localhost") + .tag(com.microsoft.kiota.http.ObservabilityOptions.class, obsOptions) + .build(); + when(mockChain.request()).thenReturn(request); + when(mockChain.proceed(any(Request.class))) + .thenAnswer( + invocation -> { + final Request req = invocation.getArgument(0); + assertSame(mockSpan, req.tag(Span.class)); + verify(mockSpan, never()).end(); + verify(mockScope, never()).close(); + when(mockResponse.request()).thenReturn(req); + return mockResponse; + }); + + final Response response = handler.intercept(mockChain); + assertNotNull(response); + verify(mockScope).close(); + verify(mockSpan).end(); + } + } + + @Test + void enclosesProceedAndCleansUpOnException() throws IOException { + try (MockedStatic otelMock = mockStatic(GlobalOpenTelemetry.class)) { + final Tracer mockTracer = mock(Tracer.class); + final SpanBuilder mockSpanBuilder = mock(SpanBuilder.class); + final Span mockSpan = mock(Span.class); + final io.opentelemetry.context.Scope mockScope = + mock(io.opentelemetry.context.Scope.class); + + otelMock.when(() -> GlobalOpenTelemetry.getTracer(anyString())).thenReturn(mockTracer); + when(mockTracer.spanBuilder(anyString())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.setParent(any())).thenReturn(mockSpanBuilder); + when(mockSpanBuilder.startSpan()).thenReturn(mockSpan); + when(mockSpan.makeCurrent()).thenReturn(mockScope); + + final com.microsoft.kiota.http.ObservabilityOptions obsOptions = + new com.microsoft.kiota.http.ObservabilityOptions(); + final UserAgentHandler handler = new UserAgentHandler(); + final Request request = + new Request.Builder() + .url("http://localhost") + .tag(com.microsoft.kiota.http.ObservabilityOptions.class, obsOptions) + .build(); + when(mockChain.request()).thenReturn(request); + when(mockChain.proceed(any(Request.class))) + .thenAnswer( + invocation -> { + final Request req = invocation.getArgument(0); + assertSame(mockSpan, req.tag(Span.class)); + verify(mockSpan, never()).end(); + verify(mockScope, never()).close(); + throw new IOException("network failure"); + }); + + assertThrows(IOException.class, () -> handler.intercept(mockChain)); + verify(mockScope).close(); + verify(mockSpan).end(); + } + } }