natsdotnet/tests/NATS.Server.Tests/Internal/TraceContextPropagationTests.cs

using System.Text;
using NATS.Server.Internal;

namespace NATS.Server.Tests.Internal;

/// <summary>
/// Tests for TraceContextPropagator: trace creation, header injection/extraction,
/// child span creation, round-trip fidelity, and ShouldTrace detection.
/// Go reference: server/msgtrace.go — trace context embedding and extraction.
/// </summary>
public class TraceContextPropagationTests
{
    // Helper: build a minimal NATS/1.0 header block with the given headers.
    private static byte[] BuildNatsHeaders(params (string key, string value)[] headers)
    {
        var sb = new StringBuilder("NATS/1.0\r\n");
        foreach (var (key, value) in headers)
            sb.Append($"{key}: {value}\r\n");
        sb.Append("\r\n");
        return Encoding.ASCII.GetBytes(sb.ToString());
    }

    [Fact]
    public void CreateTrace_GeneratesValidContext()
    {
        var ctx = TraceContextPropagator.CreateTrace("abc123", "span456", destination: "trace.dest");

        ctx.TraceId.ShouldBe("abc123");
        ctx.SpanId.ShouldBe("span456");
        ctx.Destination.ShouldBe("trace.dest");
        ctx.TraceOnly.ShouldBeFalse();
        ctx.CreatedAt.ShouldBeInRange(DateTime.UtcNow.AddSeconds(-5), DateTime.UtcNow.AddSeconds(1));
    }

    [Fact]
    public void ExtractTrace_ValidHeaders_ReturnsContext()
    {
        var headers = BuildNatsHeaders((TraceContextPropagator.TraceParentHeader, "trace1-span1"));

        var ctx = TraceContextPropagator.ExtractTrace(headers);

        ctx.ShouldNotBeNull();
        ctx.TraceId.ShouldBe("trace1");
        ctx.SpanId.ShouldBe("span1");
    }

    [Fact]
    public void ExtractTrace_NoTraceHeader_ReturnsNull()
    {
        var headers = BuildNatsHeaders(("Content-Type", "text/plain"));

        var ctx = TraceContextPropagator.ExtractTrace(headers);

        ctx.ShouldBeNull();
    }

    [Fact]
    public void InjectTrace_AppendsToHeaders()
    {
        var existing = BuildNatsHeaders(("Content-Type", "text/plain"));
        var ctx = TraceContextPropagator.CreateTrace("tid", "sid");

        var result = TraceContextPropagator.InjectTrace(ctx, existing);

        var text = Encoding.ASCII.GetString(result);
        text.ShouldContain($"{TraceContextPropagator.TraceParentHeader}: tid-sid");
        text.ShouldContain("Content-Type: text/plain");
    }

    [Fact]
    public void InjectTrace_EmptyHeaders_CreatesNew()
    {
        var ctx = TraceContextPropagator.CreateTrace("newtrace", "newspan");

        var result = TraceContextPropagator.InjectTrace(ctx, ReadOnlySpan<byte>.Empty);

        var text = Encoding.ASCII.GetString(result);
        text.ShouldStartWith("NATS/1.0\r\n");
        text.ShouldContain($"{TraceContextPropagator.TraceParentHeader}: newtrace-newspan");
    }

    [Fact]
    public void CreateChildSpan_PreservesTraceId()
    {
        var parent = TraceContextPropagator.CreateTrace("parentTrace", "parentSpan");

        var child = TraceContextPropagator.CreateChildSpan(parent, "childSpan");

        child.TraceId.ShouldBe("parentTrace");
    }

    [Fact]
    public void CreateChildSpan_NewSpanId()
    {
        var parent = TraceContextPropagator.CreateTrace("parentTrace", "parentSpan");

        var child = TraceContextPropagator.CreateChildSpan(parent, "childSpan");

        child.SpanId.ShouldBe("childSpan");
        child.SpanId.ShouldNotBe(parent.SpanId);
    }

    [Fact]
    public void ShouldTrace_WithHeader_ReturnsTrue()
    {
        var headers = BuildNatsHeaders((TraceContextPropagator.TraceParentHeader, "trace1-span1"));

        TraceContextPropagator.ShouldTrace(headers).ShouldBeTrue();
    }

    [Fact]
    public void ShouldTrace_WithoutHeader_ReturnsFalse()
    {
        var headers = BuildNatsHeaders(("Content-Type", "text/plain"));

        TraceContextPropagator.ShouldTrace(headers).ShouldBeFalse();
    }

    [Fact]
    public void RoundTrip_CreateInjectExtract_Matches()
    {
        // Use hex-style IDs (no dashes) so the "{traceId}-{spanId}" wire format
        // can be unambiguously split on the single separator dash.
        var original = TraceContextPropagator.CreateTrace("0af7651916cd43dd8448eb211c80319c", "b7ad6b7169203331", destination: "trace.dest");

        // Inject into empty headers
        var injected = TraceContextPropagator.InjectTrace(original, ReadOnlySpan<byte>.Empty);

        // Extract back from the injected headers
        var extracted = TraceContextPropagator.ExtractTrace(injected);

        extracted.ShouldNotBeNull();
        extracted.TraceId.ShouldBe(original.TraceId);
        extracted.SpanId.ShouldBe(original.SpanId);
    }
}