Files
sub2api-cn-relay-manager/internal/app/route_proxy_api_test.go
2026-05-29 10:49:27 +08:00

611 lines
21 KiB
Go

package app
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"testing"
"sub2api-cn-relay-manager/internal/store/sqlite"
)
func TestAPIProxyRouteChatCompletionsReturnsResolveAndForward(t *testing.T) {
t.Parallel()
handler := NewAPIHandler("secret-token", ActionSet{
ProxyRouteChatCompletions: func(_ context.Context, req ProxyRouteChatCompletionsRequest) (ProxyRouteChatCompletionsResult, error) {
if req.LogicalGroupID != "gpt-shared" {
t.Fatalf("LogicalGroupID = %q, want gpt-shared", req.LogicalGroupID)
}
if req.GatewayAPIKey != "gateway-key" {
t.Fatalf("GatewayAPIKey = %q, want gateway-key", req.GatewayAPIKey)
}
return ProxyRouteChatCompletionsResult{
Resolve: ResolveRouteInfo{
RequestID: "req-proxy-1",
Backend: "memory",
LogicalGroupID: req.LogicalGroupID,
PublicModel: req.PublicModel,
StickyKey: "lg:gpt-shared:m:gpt-5.4:conv:conv-1",
StickyHit: false,
StickyAction: "bind",
RouteID: "asxs",
ShadowGroupID: "gpt-shared__asxs",
ShadowHostID: "remote43",
ShadowModel: "gpt-5.4-asxs",
},
Forward: RouteChatCompletionsForwardInfo{
OK: true,
HostID: "remote43",
HostBaseURL: "https://sub2api.example.com",
ShadowGroupID: "gpt-shared__asxs",
ShadowModel: "gpt-5.4-asxs",
UpstreamPath: "/v1/chat/completions",
UpstreamStatus: 200,
LatencyMS: 12,
ContentType: "application/json",
Response: map[string]any{
"id": "chatcmpl_proxy",
},
},
}, nil
},
})
request := httptestRequest(t, http.MethodPost, "/api/routing/proxy/chat/completions", map[string]any{
"logical_group_id": "gpt-shared",
"public_model": "gpt-5.4",
"scope": "conversation",
"subject_id": "conv-1",
"gateway_api_key": "gateway-key",
"sync": true,
}, "secret-token")
response := httptestRecorder(handler, request)
assertStatusCode(t, response, http.StatusOK)
assertJSONContains(t, response.Body().Bytes(), "resolve.route_id", "asxs")
assertJSONContains(t, response.Body().Bytes(), "forward.shadow_model", "gpt-5.4-asxs")
assertJSONContains(t, response.Body().Bytes(), "forward.upstream_status", float64(200))
}
func TestNewActionSetProxyRouteChatCompletionsFlow(t *testing.T) {
t.Parallel()
var (
gotAuthHeader string
gotModel string
gotPrompt string
)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/chat/completions" {
t.Fatalf("URL.Path = %q, want /v1/chat/completions", r.URL.Path)
}
gotAuthHeader = r.Header.Get("Authorization")
var payload struct {
Model string `json:"model"`
Messages []struct {
Role string `json:"role"`
Content string `json:"content"`
} `json:"messages"`
}
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
t.Fatalf("json.Decode() error = %v", err)
}
gotModel = payload.Model
if len(payload.Messages) > 0 {
gotPrompt = payload.Messages[0].Content
}
writeJSON(w, http.StatusOK, map[string]any{
"id": "chatcmpl_proxy",
"object": "chat.completion",
"choices": []map[string]any{
{
"index": 0,
"message": map[string]any{
"role": "assistant",
"content": "pong",
},
},
},
})
}))
defer server.Close()
dsn := "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "route-proxy.db")) + "?_busy_timeout=5000"
actions := NewActionSet(dsn)
ctx := context.Background()
store, err := sqlite.Open(ctx, dsn)
if err != nil {
t.Fatalf("sqlite.Open() error = %v", err)
}
defer store.Close()
if _, err := store.Hosts().Create(ctx, sqlite.Host{
HostID: "remote43",
BaseURL: server.URL,
HostVersion: "0.1.126",
AuthType: "apikey",
AuthToken: "host-admin-token",
}); err != nil {
t.Fatalf("Hosts().Create() error = %v", err)
}
if _, err := actions.CreateLogicalGroup(ctx, CreateLogicalGroupRequest{
LogicalGroupID: "gpt-shared",
DisplayName: "GPT Shared",
Status: "active",
RoutePolicy: "priority",
StickyMode: "conversation_preferred",
ConversationTTLSeconds: 1200,
UserModelTTLSeconds: 600,
FailoverThreshold: 2,
CooldownSeconds: 300,
}); err != nil {
t.Fatalf("CreateLogicalGroup() error = %v", err)
}
if _, err := actions.CreateLogicalGroupModel(ctx, CreateLogicalGroupModelRequest{
LogicalGroupID: "gpt-shared",
PublicModel: "gpt-5.4",
Status: "active",
}); err != nil {
t.Fatalf("CreateLogicalGroupModel() error = %v", err)
}
if _, err := actions.CreateLogicalGroupRoute(ctx, CreateLogicalGroupRouteRequest{
LogicalGroupID: "gpt-shared",
RouteID: "asxs",
Name: "ASXS",
Status: "active",
Priority: 10,
ShadowGroupID: "gpt-shared__asxs",
ShadowHostID: "remote43",
UpstreamBaseURLHint: "https://api.asxs.top/v1",
}); err != nil {
t.Fatalf("CreateLogicalGroupRoute() error = %v", err)
}
if _, err := actions.CreateLogicalGroupRouteModel(ctx, CreateLogicalGroupRouteModelRequest{
LogicalGroupID: "gpt-shared",
RouteID: "asxs",
PublicModel: "gpt-5.4",
ShadowModel: "gpt-5.4-asxs",
Status: "active",
}); err != nil {
t.Fatalf("CreateLogicalGroupRouteModel() error = %v", err)
}
result, err := actions.ProxyRouteChatCompletions(ctx, ProxyRouteChatCompletionsRequest{
RequestID: "req-proxy-1",
LogicalGroupID: "gpt-shared",
PublicModel: "gpt-5.4",
Scope: "conversation",
SubjectID: "conv-1",
GatewayAPIKey: "gateway-key",
Sync: true,
})
if err != nil {
t.Fatalf("ProxyRouteChatCompletions() error = %v", err)
}
if gotAuthHeader != "Bearer gateway-key" {
t.Fatalf("Authorization header = %q, want Bearer gateway-key", gotAuthHeader)
}
if gotModel != "gpt-5.4-asxs" {
t.Fatalf("forwarded model = %q, want gpt-5.4-asxs", gotModel)
}
if gotPrompt != "ping" {
t.Fatalf("forwarded prompt = %q, want ping", gotPrompt)
}
if result.Resolve.RouteID != "asxs" || result.Resolve.ShadowModel != "gpt-5.4-asxs" {
t.Fatalf("Resolve = %+v, want selected asxs route with shadow model", result.Resolve)
}
if !result.Forward.OK || result.Forward.UpstreamStatus != http.StatusOK {
t.Fatalf("Forward = %+v, want successful 200 forward", result.Forward)
}
if result.Forward.HostID != "remote43" || result.Forward.HostBaseURL != server.URL {
t.Fatalf("Forward = %+v, want host remote43 and server URL", result.Forward)
}
decisions, err := actions.ListRouteDecisionLogs(ctx, ListRouteDecisionLogsRequest{
RequestID: "req-proxy-1",
Limit: 10,
})
if err != nil {
t.Fatalf("ListRouteDecisionLogs() error = %v", err)
}
if len(decisions) != 2 {
t.Fatalf("ListRouteDecisionLogs() len = %d, want 2", len(decisions))
}
if decisions[0].UpstreamStatus != http.StatusOK || decisions[0].SelectedRouteID != "asxs" {
t.Fatalf("latest decision log = %+v, want upstream_status 200 on asxs", decisions[0])
}
if decisions[1].UpstreamStatus != 0 {
t.Fatalf("initial decision log = %+v, want upstream_status 0 before forward", decisions[1])
}
}
func TestNewActionSetProxyRouteChatCompletionsManagedSubscriptionFlow(t *testing.T) {
t.Parallel()
var (
gotAuthHeader string
gotModel string
)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch {
case r.Method == http.MethodGet && strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/users?"):
_, _ = w.Write([]byte(`{"data":{"items":[]}}`))
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/users":
_, _ = w.Write([]byte(`{"data":{"id":84,"email":"relay-sub-managed-user@sub2api.local"}}`))
case r.Method == http.MethodPut && r.URL.Path == "/api/v1/admin/users/84":
_, _ = w.Write([]byte(`{"data":{"id":84}}`))
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/users/84/balance":
_, _ = w.Write([]byte(`{"data":{"id":84}}`))
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/subscriptions/assign":
_, _ = w.Write([]byte(`{"data":{"id":401}}`))
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/auth/login":
_, _ = w.Write([]byte(`{"data":{"access_token":"user-jwt"}}`))
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/keys":
_, _ = w.Write([]byte(`{"data":{"id":501,"key":"sk-relay-key","name":"managed-key"}}`))
case r.Method == http.MethodPut && r.URL.Path == "/api/v1/admin/api-keys/501":
_, _ = w.Write([]byte(`{"data":{"api_key":{"id":501}}}`))
case r.Method == http.MethodPost && r.URL.Path == "/v1/chat/completions":
gotAuthHeader = r.Header.Get("Authorization")
var payload struct {
Model string `json:"model"`
}
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
t.Fatalf("json.Decode() error = %v", err)
}
gotModel = payload.Model
writeJSON(w, http.StatusOK, map[string]any{
"id": "chatcmpl_proxy_managed",
"choices": []map[string]any{
{
"message": map[string]any{
"role": "assistant",
"content": "pong-managed",
},
},
},
})
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer server.Close()
dsn := "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "route-proxy-managed.db")) + "?_busy_timeout=5000"
actions := NewActionSet(dsn)
ctx := context.Background()
store, err := sqlite.Open(ctx, dsn)
if err != nil {
t.Fatalf("sqlite.Open() error = %v", err)
}
defer store.Close()
hostID, err := store.Hosts().Create(ctx, sqlite.Host{
HostID: "remote43-managed",
BaseURL: server.URL,
HostVersion: "0.1.126",
AuthType: "bearer",
AuthToken: "host-admin-token",
})
if err != nil {
t.Fatalf("Hosts().Create() error = %v", err)
}
packID, err := store.Packs().Create(ctx, sqlite.Pack{
PackID: "managed-pack",
Version: "1.0.0",
Checksum: "sha256-managed-pack",
Vendor: "tksea",
ManifestJSON: "{}",
})
if err != nil {
t.Fatalf("Packs().Create() error = %v", err)
}
providerID, err := store.Providers().Create(ctx, sqlite.Provider{
PackID: packID,
ProviderID: "managed-provider",
DisplayName: "Managed Provider",
BaseURL: "https://api.asxs.top/v1",
Platform: "openai",
})
if err != nil {
t.Fatalf("Providers().Create() error = %v", err)
}
batchID, err := store.ImportBatches().Create(ctx, sqlite.ImportBatch{
HostID: hostID,
PackID: packID,
ProviderID: providerID,
Mode: "strict",
BatchStatus: "succeeded",
AccessStatus: "subscription_ready",
})
if err != nil {
t.Fatalf("ImportBatches().Create() error = %v", err)
}
if _, err := actions.CreateLogicalGroup(ctx, CreateLogicalGroupRequest{
LogicalGroupID: "gpt-shared-managed",
DisplayName: "GPT Shared Managed",
Status: "active",
RoutePolicy: "priority",
StickyMode: "conversation_preferred",
ConversationTTLSeconds: 1200,
UserModelTTLSeconds: 600,
FailoverThreshold: 2,
CooldownSeconds: 300,
}); err != nil {
t.Fatalf("CreateLogicalGroup() error = %v", err)
}
if _, err := actions.CreateLogicalGroupModel(ctx, CreateLogicalGroupModelRequest{
LogicalGroupID: "gpt-shared-managed",
PublicModel: "gpt-5.4",
Status: "active",
}); err != nil {
t.Fatalf("CreateLogicalGroupModel() error = %v", err)
}
if _, err := actions.CreateLogicalGroupRoute(ctx, CreateLogicalGroupRouteRequest{
LogicalGroupID: "gpt-shared-managed",
RouteID: "asxs-managed",
Name: "ASXS Managed",
Status: "active",
Priority: 10,
ShadowGroupID: "101",
ShadowHostID: "remote43-managed",
UpstreamBaseURLHint: "https://api.asxs.top/v1",
}); err != nil {
t.Fatalf("CreateLogicalGroupRoute() error = %v", err)
}
if _, err := actions.CreateLogicalGroupRouteModel(ctx, CreateLogicalGroupRouteModelRequest{
LogicalGroupID: "gpt-shared-managed",
RouteID: "asxs-managed",
PublicModel: "gpt-5.4",
ShadowModel: "gpt-5.4-asxs",
Status: "active",
}); err != nil {
t.Fatalf("CreateLogicalGroupRouteModel() error = %v", err)
}
if _, err := store.ManagedResources().Create(ctx, sqlite.ManagedResource{
BatchID: batchID,
HostID: hostID,
ResourceType: "group",
HostResourceID: "101",
ResourceName: "shadow-group-asxs",
}); err != nil {
t.Fatalf("ManagedResources().Create() error = %v", err)
}
result, err := actions.ProxyRouteChatCompletions(ctx, ProxyRouteChatCompletionsRequest{
RequestID: "req-proxy-managed-1",
LogicalGroupID: "gpt-shared-managed",
PublicModel: "gpt-5.4",
Scope: "conversation",
SubjectID: "conv-managed-1",
SubscriptionUserID: "crm-user-1",
Sync: true,
})
if err != nil {
t.Fatalf("ProxyRouteChatCompletions() error = %v", err)
}
if !strings.HasPrefix(gotAuthHeader, "Bearer sk-relay-") {
t.Fatalf("Authorization header = %q, want Bearer sk-relay-*", gotAuthHeader)
}
if gotModel != "gpt-5.4-asxs" {
t.Fatalf("forwarded model = %q, want gpt-5.4-asxs", gotModel)
}
if result.Forward.EffectiveGatewayKeySource != "managed_subscription" {
t.Fatalf("EffectiveGatewayKeySource = %q, want managed_subscription", result.Forward.EffectiveGatewayKeySource)
}
if result.Forward.EffectiveGatewayKeyFingerprint == "" {
t.Fatal("EffectiveGatewayKeyFingerprint = empty, want hashed managed key fingerprint")
}
if result.Forward.ManagedUserID != "84" {
t.Fatalf("ManagedUserID = %q, want 84", result.Forward.ManagedUserID)
}
}
func TestProxyChatCompletionToShadowHostReportsNon2xx(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusTooManyRequests, map[string]any{
"error": map[string]any{
"message": "rate limited",
},
})
}))
defer server.Close()
info := proxyChatCompletionToShadowHost(context.Background(), server.URL, "gateway-key", "gpt-5.4-asxs", nil, 0, nil)
if info.OK {
t.Fatalf("proxyChatCompletionToShadowHost() = %+v, want non-ok result", info)
}
if info.UpstreamStatus != http.StatusTooManyRequests || info.ErrorClass != "gateway_rate_limited" {
t.Fatalf("proxyChatCompletionToShadowHost() = %+v, want 429 gateway_rate_limited", info)
}
response, ok := info.Response.(map[string]any)
if !ok || response["error"] == nil {
t.Fatalf("proxyChatCompletionToShadowHost() response = %#v, want decoded json body", info.Response)
}
}
func TestRouteProxyHelpers(t *testing.T) {
t.Parallel()
if got := normalizeProxyMaxTokens(0); got != 8 {
t.Fatalf("normalizeProxyMaxTokens(0) = %d, want 8", got)
}
if got := normalizeProxyTemperature(nil); got != 0 {
t.Fatalf("normalizeProxyTemperature(nil) = %v, want 0", got)
}
if got := normalizeProxyChatMessages(nil); len(got) != 1 || got[0]["content"] != "ping" {
t.Fatalf("normalizeProxyChatMessages(nil) = %#v, want default ping message", got)
}
if got := classifyProxyUpstreamStatus(http.StatusForbidden); got != "gateway_auth_error" {
t.Fatalf("classifyProxyUpstreamStatus(403) = %q, want gateway_auth_error", got)
}
if _, err := joinRouteProxyPath("://bad-url", routeChatCompletionsPath); err == nil {
t.Fatal("joinRouteProxyPath(invalid) error = nil, want error")
}
if got := resolveProxyUserKey(ProxyRouteChatCompletionsRequest{Scope: "user", SubjectID: "user-1"}); got != "user-1" {
t.Fatalf("resolveProxyUserKey(user) = %q, want user-1", got)
}
if got := resolveProxyConversationKey(ProxyRouteChatCompletionsRequest{Scope: "conversation", SubjectID: "conv-1"}); got != "conv-1" {
t.Fatalf("resolveProxyConversationKey(conversation) = %q, want conv-1", got)
}
}
func TestResolveShadowGroupHostResourceID(t *testing.T) {
t.Parallel()
dsn := "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "route-proxy-helper.db")) + "?_busy_timeout=5000"
ctx := context.Background()
store, err := sqlite.Open(ctx, dsn)
if err != nil {
t.Fatalf("sqlite.Open() error = %v", err)
}
defer store.Close()
hostID, err := store.Hosts().Create(ctx, sqlite.Host{
HostID: "helper-host",
BaseURL: "https://helper.example.com",
HostVersion: "0.1.126",
AuthType: "bearer",
AuthToken: "host-token",
})
if err != nil {
t.Fatalf("Hosts().Create() error = %v", err)
}
hostRow, err := store.Hosts().GetByID(ctx, hostID)
if err != nil {
t.Fatalf("Hosts().GetByID() error = %v", err)
}
if got, err := resolveShadowGroupHostResourceID(ctx, store, hostRow, nil, "101"); err != nil || got != "101" {
t.Fatalf("resolveShadowGroupHostResourceID(numeric) = (%q, %v), want 101", got, err)
}
packID, err := store.Packs().Create(ctx, sqlite.Pack{
PackID: "helper-pack",
Version: "1.0.0",
Checksum: "sha256-helper",
Vendor: "tksea",
ManifestJSON: "{}",
})
if err != nil {
t.Fatalf("Packs().Create() error = %v", err)
}
providerID, err := store.Providers().Create(ctx, sqlite.Provider{
PackID: packID,
ProviderID: "helper-provider",
DisplayName: "Helper Provider",
BaseURL: "https://helper.example.com/v1",
Platform: "openai",
})
if err != nil {
t.Fatalf("Providers().Create() error = %v", err)
}
batchID, err := store.ImportBatches().Create(ctx, sqlite.ImportBatch{
HostID: hostID,
PackID: packID,
ProviderID: providerID,
Mode: "strict",
BatchStatus: "succeeded",
AccessStatus: "subscription_ready",
})
if err != nil {
t.Fatalf("ImportBatches().Create() error = %v", err)
}
if _, err := store.ManagedResources().Create(ctx, sqlite.ManagedResource{
BatchID: batchID,
HostID: hostID,
ResourceType: "group",
HostResourceID: "202",
ResourceName: "shadow-group-name",
}); err != nil {
t.Fatalf("ManagedResources().Create() error = %v", err)
}
if got, err := resolveShadowGroupHostResourceID(ctx, store, hostRow, nil, "202"); err != nil || got != "202" {
t.Fatalf("resolveShadowGroupHostResourceID(store identity) = (%q, %v), want 202", got, err)
}
}
func TestResolveShadowGroupHostResourceIDFallsBackToHostList(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch {
case strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/groups"):
_, _ = w.Write([]byte(`{"data":[{"id":"303","name":"shadow-group-remote"}]}`))
case strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/channels"):
_, _ = w.Write([]byte(`{"data":[]}`))
case strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/payment/plans"):
_, _ = w.Write([]byte(`{"data":[]}`))
case strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/accounts"):
_, _ = w.Write([]byte(`{"data":{"items":[],"pages":1}}`))
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer server.Close()
dsn := "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "route-proxy-fallback.db")) + "?_busy_timeout=5000"
ctx := context.Background()
store, err := sqlite.Open(ctx, dsn)
if err != nil {
t.Fatalf("sqlite.Open() error = %v", err)
}
defer store.Close()
hostID, err := store.Hosts().Create(ctx, sqlite.Host{
HostID: "fallback-host",
BaseURL: server.URL,
HostVersion: "0.1.126",
AuthType: "bearer",
AuthToken: "host-token",
})
if err != nil {
t.Fatalf("Hosts().Create() error = %v", err)
}
hostRow, err := store.Hosts().GetByID(ctx, hostID)
if err != nil {
t.Fatalf("Hosts().GetByID() error = %v", err)
}
hostClient, err := newSub2APIClient(server.URL, CreateHostAuth{Type: "bearer", Token: "host-token"})
if err != nil {
t.Fatalf("newSub2APIClient() error = %v", err)
}
got, err := resolveShadowGroupHostResourceID(ctx, store, hostRow, hostClient, "shadow-group-remote")
if err != nil {
t.Fatalf("resolveShadowGroupHostResourceID(host fallback) error = %v", err)
}
if got != "303" {
t.Fatalf("resolveShadowGroupHostResourceID(host fallback) = %q, want 303", got)
}
}
func TestAPIProxyRouteChatCompletionsRejectsMissingGatewayAndSubscriptionUser(t *testing.T) {
t.Parallel()
handler := NewAPIHandler("secret-token", ActionSet{
ProxyRouteChatCompletions: buildProxyRouteChatCompletionsAction("file::memory:?cache=shared", func(context.Context, ResolveRouteRequest) (ResolveRouteInfo, error) {
t.Fatal("ResolveRoute should not be called when auth inputs are missing")
return ResolveRouteInfo{}, nil
}, newLazyRouteLogWriter("file::memory:?cache=shared")),
})
request := httptestRequest(t, http.MethodPost, "/api/routing/proxy/chat/completions", map[string]any{
"logical_group_id": "gpt-shared",
"public_model": "gpt-5.4",
"scope": "conversation",
"subject_id": "conv-1",
}, "secret-token")
response := httptestRecorder(handler, request)
assertStatusCode(t, response, http.StatusBadRequest)
assertJSONContains(t, response.Body().Bytes(), "error.message", "gateway_api_key or subscription_user_id is required")
}