using System.Net; using System.Net.Http; using System.Net.Http.Headers; using Omika.Core.Observability.Diagnostic; namespace Omika.Tests.Observability.Correlation; public class TraceIdPropagationHandlerTests { private sealed class FakeInnerHandler : HttpMessageHandler { public HttpRequestMessage? LastRequest { get; private set; } protected override Task SendAsync( HttpRequestMessage request, CancellationToken cancellationToken) { LastRequest = request; return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)); } } [Fact] public async Task SendAsync_NoCurrentTrace_GeneratesTraceIdHeader() { var traceContext = new TraceContext(); var inner = new FakeInnerHandler(); var handler = new TraceIdPropagationHandler(traceContext, innerHandler: inner); var client = new HttpClient(handler); await client.GetAsync("http://test.local/"); Assert.NotNull(inner.LastRequest); Assert.True(inner.LastRequest.Headers.TryGetValues(CorrelationHeaders.TraceId, out var values)); Assert.True(Guid.TryParse(Assert.Single(values), out _)); } [Fact] public async Task SendAsync_CurrentTrace_UsesIt() { var traceContext = new TraceContext(); Guid expected = Guid.NewGuid(); var inner = new FakeInnerHandler(); using (traceContext.BeginTrace(expected)) { var handler = new TraceIdPropagationHandler(traceContext, innerHandler: inner); var client = new HttpClient(handler); await client.GetAsync("http://test.local/"); } Assert.Equal( expected.ToString("D"), inner.LastRequest!.Headers.GetValues(CorrelationHeaders.TraceId).Single()); } [Fact] public async Task SendAsync_WithSession_AddsSessionHeader() { var traceContext = new TraceContext(); var session = new DiagnosticSession(); var inner = new FakeInnerHandler(); var handler = new TraceIdPropagationHandler(traceContext, session, inner); var client = new HttpClient(handler); await client.GetAsync("http://test.local/"); Assert.Equal( session.SessionId.ToString("D"), inner.LastRequest!.Headers.GetValues(CorrelationHeaders.SessionId).Single()); } [Fact] public async Task SendAsync_DoesNotModifyAuthorization() { var traceContext = new TraceContext(); var inner = new FakeInnerHandler(); var handler = new TraceIdPropagationHandler(traceContext, innerHandler: inner); var client = new HttpClient(handler); var request = new HttpRequestMessage(HttpMethod.Get, "http://test.local/"); request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", "token-abc"); await client.SendAsync(request); Assert.NotNull(inner.LastRequest); Assert.Equal("Bearer", inner.LastRequest.Headers.Authorization!.Scheme); Assert.Equal("token-abc", inner.LastRequest.Headers.Authorization!.Parameter); } [Fact] public async Task SendAsync_ExistingTraceHeader_NotOverwritten() { var traceContext = new TraceContext(); var inner = new FakeInnerHandler(); var handler = new TraceIdPropagationHandler(traceContext, innerHandler: inner); var client = new HttpClient(handler); Guid existing = Guid.NewGuid(); var request = new HttpRequestMessage(HttpMethod.Get, "http://test.local/"); request.Headers.Add(CorrelationHeaders.TraceId, existing.ToString("D")); await client.SendAsync(request); Assert.Equal( existing.ToString("D"), inner.LastRequest!.Headers.GetValues(CorrelationHeaders.TraceId).Single()); } [Fact] public async Task SendAsync_DoesNotModifyPayload() { var traceContext = new TraceContext(); var inner = new FakeInnerHandler(); var handler = new TraceIdPropagationHandler(traceContext, innerHandler: inner); var client = new HttpClient(handler); var request = new HttpRequestMessage(HttpMethod.Post, "http://test.local/"); request.Content = new StringContent("{\"a\":1}", System.Text.Encoding.UTF8, "application/json"); await client.SendAsync(request); Assert.NotNull(inner.LastRequest); string body = await inner.LastRequest!.Content!.ReadAsStringAsync(); Assert.Equal("{\"a\":1}", body); } }