Omika/tests/Omika.Tests/Observability/Correlation/TraceIdPropagationHandlerTests.cs

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<HttpResponseMessage> 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);
    }
}