mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-05 08:55:28 +08:00
48 lines
1.3 KiB
Go
48 lines
1.3 KiB
Go
package openai
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/go-chi/chi/v5/middleware"
|
|
)
|
|
|
|
func traceIDViaMiddleware(req *http.Request) string {
|
|
if req == nil {
|
|
return requestTraceID(nil)
|
|
}
|
|
var got string
|
|
h := middleware.RequestID(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
|
|
got = requestTraceID(r)
|
|
}))
|
|
h.ServeHTTP(httptest.NewRecorder(), req)
|
|
return got
|
|
}
|
|
|
|
func TestRequestTraceIDPriority(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions?__trace_id=query-trace", nil)
|
|
req.Header.Set("X-Ds2-Test-Trace", "header-trace")
|
|
got := traceIDViaMiddleware(req)
|
|
if got != "query-trace" {
|
|
t.Fatalf("expected query trace id to win, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestRequestTraceIDHeaderFallback(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)
|
|
req.Header.Set("X-Ds2-Test-Trace", "header-trace")
|
|
got := traceIDViaMiddleware(req)
|
|
if got != "header-trace" {
|
|
t.Fatalf("expected header trace id to win when query missing, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestRequestTraceIDReqIDFallback(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)
|
|
got := traceIDViaMiddleware(req)
|
|
if got == "" {
|
|
t.Fatal("expected middleware request id fallback to be non-empty")
|
|
}
|
|
}
|