sub2api/backend/internal/pkg/servertiming/http_test.go

package servertiming

import (
	"context"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"
)

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}

type trackingBody struct {
	read bool
}

func (b *trackingBody) Read(_ []byte) (int, error) {
	b.read = true
	return 0, io.EOF
}

func (b *trackingBody) Close() error { return nil }

func TestWrapRoundTripperRecordsResponseHeaderLatency(t *testing.T) {
	startedAt := time.Now()
	collector := New(startedAt)
	body := &trackingBody{}
	baseCalled := false
	base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		baseCalled = true
		return &http.Response{
			StatusCode: http.StatusOK,
			Body:       body,
			Header:     make(http.Header),
			Request:    req,
		}, nil
	})
	req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.github.com/repos/example/project", nil)
	if err != nil {
		t.Fatal(err)
	}

	resp, err := WrapRoundTripper(base).RoundTrip(req)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = resp.Body.Close() }()
	if !baseCalled {
		t.Fatal("base RoundTripper was not called")
	}
	if body.read {
		t.Fatal("RoundTripper instrumentation read the response body; timing must stop at response headers")
	}
	header := collector.HeaderValue(time.Now(), "bypass")
	if !strings.Contains(header, `dep_github;dur=`) || !strings.Contains(header, `deps;dur=`) {
		t.Fatalf("dependency metrics missing from header: %q", header)
	}
}

func TestWrapRoundTripperUsesContextModuleOverride(t *testing.T) {
	collector := New(time.Now())
	ctx := WithDependencyModule(WithCollector(context.Background(), collector), "data-managementd")
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://private.example.test/path", nil)
	if err != nil {
		t.Fatal(err)
	}
	base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
	})

	if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
		t.Fatal(err)
	}
	header := collector.HeaderValue(time.Now(), "bypass")
	if !strings.Contains(header, "dep_data_managementd") {
		t.Fatalf("module override missing from header: %q", header)
	}
	if strings.Contains(header, "private.example") {
		t.Fatalf("raw host leaked into header: %q", header)
	}
}

func TestWrapRoundTripperSkipsInactiveContext(t *testing.T) {
	called := false
	base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		called = true
		return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
	})
	req, err := http.NewRequest(http.MethodGet, "https://api.openai.com/v1/models", nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
		t.Fatal(err)
	}
	if !called {
		t.Fatal("inactive request did not reach base RoundTripper")
	}
}

func TestDoRecordsWithoutChangingTransportType(t *testing.T) {
	base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
	})
	client := &http.Client{Transport: base}
	collector := New(time.Now())
	req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.openai.com/v1/models", nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := Do(client, req); err != nil {
		t.Fatal(err)
	}
	if _, ok := client.Transport.(roundTripFunc); !ok {
		t.Fatalf("Do changed client transport type to %T", client.Transport)
	}
	if header := collector.HeaderValue(time.Now(), "bypass"); !strings.Contains(header, "dep_openai;dur=") {
		t.Fatalf("dependency metric missing from header: %q", header)
	}
}

func TestDependencyModuleClassification(t *testing.T) {
	tests := map[string]string{
		"https://api.github.com/repos/a/b":                    "github",
		"https://api.openai.com/v1/models":                    "openai",
		"https://api.anthropic.com/v1/messages":               "anthropic",
		"https://generativelanguage.googleapis.com/v1/models": "gemini",
		"https://cloudcode-pa.googleapis.com/v1internal":      "antigravity",
		"https://storage.googleapis.com/bucket/object":        "google",
		"https://bucket.s3.amazonaws.com/object":              "s3",
		"https://api.stripe.com/v1/refunds":                   "payment",
		"https://dependency.example.test/path":                "http",
	}
	for rawURL, want := range tests {
		req, err := http.NewRequest(http.MethodGet, rawURL, nil)
		if err != nil {
			t.Fatalf("NewRequest(%q): %v", rawURL, err)
		}
		if got := dependencyModule(req); got != want {
			t.Errorf("dependencyModule(%q) = %q, want %q", rawURL, got, want)
		}
	}
}

func TestClientInstrumentationDoesNotMutateOriginal(t *testing.T) {
	base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
	})
	original := &http.Client{Transport: base, Timeout: time.Second}
	instrumented := InstrumentClient(original)
	if instrumented == original {
		t.Fatal("InstrumentClient returned the original client")
	}
	if _, ok := original.Transport.(roundTripFunc); !ok {
		t.Fatalf("InstrumentClient mutated the original transport to %T", original.Transport)
	}
	if instrumented.Timeout != original.Timeout {
		t.Fatal("InstrumentClient did not preserve client settings")
	}
	if WrapRoundTripper(instrumented.Transport) != instrumented.Transport {
		t.Fatal("WrapRoundTripper wrapped an already instrumented transport twice")
	}
}