277 lines
9.2 KiB
Go
277 lines
9.2 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"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 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)
|
|
}
|
|
}
|