Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/pkg/servertiming/http_test.go

@@ -0,0 +1,168 @@
+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")
+	}
+}