Files
2026-07-13 12:35:43 +08:00

291 lines
9.4 KiB
Go

package auth
import (
"context"
"testing"
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
)
func TestLookupAPIKeyUpstreamModel(t *testing.T) {
cfg := &internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{
{
APIKey: "k",
BaseURL: "https://example.com",
Models: []internalconfig.GeminiModel{
{Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"},
{Name: "gemini-2.5-flash(low)", Alias: "g25f"},
},
},
},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
_, _ = mgr.Register(ctx, &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k", "base_url": "https://example.com"}})
tests := []struct {
name string
authID string
input string
want string
}{
// Fast path + suffix preservation
{"alias with suffix", "a1", "g25p(8192)", "gemini-2.5-pro-exp-03-25(8192)"},
{"alias without suffix", "a1", "g25p", "gemini-2.5-pro-exp-03-25"},
// Config suffix takes priority
{"config suffix priority", "a1", "g25f(high)", "gemini-2.5-flash(low)"},
{"config suffix no user suffix", "a1", "g25f", "gemini-2.5-flash(low)"},
// Case insensitive
{"uppercase alias", "a1", "G25P", "gemini-2.5-pro-exp-03-25"},
{"mixed case with suffix", "a1", "G25p(4096)", "gemini-2.5-pro-exp-03-25(4096)"},
// Direct name lookup
{"upstream name direct", "a1", "gemini-2.5-pro-exp-03-25", "gemini-2.5-pro-exp-03-25"},
{"upstream name with suffix", "a1", "gemini-2.5-pro-exp-03-25(8192)", "gemini-2.5-pro-exp-03-25(8192)"},
// Cache miss scenarios
{"non-existent auth", "non-existent", "g25p", ""},
{"unknown alias", "a1", "unknown-alias", ""},
{"empty auth ID", "", "g25p", ""},
{"empty model", "a1", "", ""},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resolved := mgr.lookupAPIKeyUpstreamModel(tt.authID, tt.input)
if resolved != tt.want {
t.Errorf("lookupAPIKeyUpstreamModel(%q, %q) = %q, want %q", tt.authID, tt.input, resolved, tt.want)
}
})
}
}
func TestLookupAPIKeyUpstreamModel_InteractionsKey(t *testing.T) {
cfg := &internalconfig.Config{
InteractionsKey: []internalconfig.GeminiKey{{
APIKey: "interactions-key",
BaseURL: "https://interactions.example.com",
Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-flash", Alias: "native-flash"}},
}},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
_, _ = mgr.Register(ctx, &Auth{ID: "interactions-auth", Provider: "gemini-interactions", Attributes: map[string]string{"api_key": "interactions-key", "base_url": "https://interactions.example.com"}})
resolved := mgr.lookupAPIKeyUpstreamModel("interactions-auth", "native-flash")
if resolved != "gemini-2.5-flash" {
t.Fatalf("lookupAPIKeyUpstreamModel() = %q, want gemini-2.5-flash", resolved)
}
}
func TestAPIKeyModelAlias_ConfigHotReload(t *testing.T) {
cfg := &internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{
{
APIKey: "k",
Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"}},
},
},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
_, _ = mgr.Register(ctx, &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}})
// Initial alias
if resolved := mgr.lookupAPIKeyUpstreamModel("a1", "g25p"); resolved != "gemini-2.5-pro-exp-03-25" {
t.Fatalf("before reload: got %q, want %q", resolved, "gemini-2.5-pro-exp-03-25")
}
// Hot reload with new alias
mgr.SetConfig(&internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{
{
APIKey: "k",
Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-flash", Alias: "g25p"}},
},
},
})
// New alias should take effect
if resolved := mgr.lookupAPIKeyUpstreamModel("a1", "g25p"); resolved != "gemini-2.5-flash" {
t.Fatalf("after reload: got %q, want %q", resolved, "gemini-2.5-flash")
}
}
func TestAPIKeyModelAlias_MultipleProviders(t *testing.T) {
cfg := &internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{{APIKey: "gemini-key", Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro", Alias: "gp"}}}},
ClaudeKey: []internalconfig.ClaudeKey{{APIKey: "claude-key", Models: []internalconfig.ClaudeModel{{Name: "claude-sonnet-4", Alias: "cs4"}}}},
CodexKey: []internalconfig.CodexKey{{APIKey: "codex-key", Models: []internalconfig.CodexModel{{Name: "o3", Alias: "o"}}}},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
_, _ = mgr.Register(ctx, &Auth{ID: "gemini-auth", Provider: "gemini", Attributes: map[string]string{"api_key": "gemini-key"}})
_, _ = mgr.Register(ctx, &Auth{ID: "claude-auth", Provider: "claude", Attributes: map[string]string{"api_key": "claude-key"}})
_, _ = mgr.Register(ctx, &Auth{ID: "codex-auth", Provider: "codex", Attributes: map[string]string{"api_key": "codex-key"}})
tests := []struct {
authID, input, want string
}{
{"gemini-auth", "gp", "gemini-2.5-pro"},
{"claude-auth", "cs4", "claude-sonnet-4"},
{"codex-auth", "o", "o3"},
}
for _, tt := range tests {
if resolved := mgr.lookupAPIKeyUpstreamModel(tt.authID, tt.input); resolved != tt.want {
t.Errorf("lookupAPIKeyUpstreamModel(%q, %q) = %q, want %q", tt.authID, tt.input, resolved, tt.want)
}
}
}
func TestApplyAPIKeyModelAlias(t *testing.T) {
cfg := &internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{
{APIKey: "k", Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"}}},
},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
apiKeyAuth := &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}}
oauthAuth := &Auth{ID: "oauth-auth", Provider: "claude", Attributes: map[string]string{"auth_kind": "oauth"}}
_, _ = mgr.Register(ctx, apiKeyAuth)
tests := []struct {
name string
auth *Auth
inputModel string
wantModel string
}{
{
name: "api_key auth with alias",
auth: apiKeyAuth,
inputModel: "g25p(8192)",
wantModel: "gemini-2.5-pro-exp-03-25(8192)",
},
{
name: "oauth auth passthrough",
auth: oauthAuth,
inputModel: "some-model",
wantModel: "some-model",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resolvedModel := mgr.applyAPIKeyModelAlias(tt.auth, tt.inputModel)
if resolvedModel != tt.wantModel {
t.Errorf("model = %q, want %q", resolvedModel, tt.wantModel)
}
})
}
}
func TestResolveAPIKeyModelAliasWithResult_ForceMapping(t *testing.T) {
cfg := &internalconfig.Config{
ClaudeKey: []internalconfig.ClaudeKey{{
APIKey: "claude-key",
Models: []internalconfig.ClaudeModel{{
Name: "glm-5.2",
Alias: "claude-sonnet-latest",
ForceMapping: true,
}},
}},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
auth := &Auth{ID: "claude-auth", Provider: "claude", Attributes: map[string]string{"api_key": "claude-key"}}
if _, err := mgr.Register(ctx, auth); err != nil {
t.Fatalf("register auth: %v", err)
}
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "claude-sonnet-latest")
if result.UpstreamModel != "glm-5.2" || !result.ForceMapping || result.OriginalAlias != "claude-sonnet-latest" {
t.Fatalf("resolveAPIKeyModelAliasWithResult() = %+v, want upstream glm-5.2 with force mapping", result)
}
noRewrite := mgr.resolveAPIKeyModelAliasWithResult(auth, "glm-5.2")
if noRewrite.UpstreamModel != "glm-5.2" || noRewrite.ForceMapping || noRewrite.OriginalAlias != "" {
t.Fatalf("resolveAPIKeyModelAliasWithResult() direct upstream = %+v, want passthrough without rewrite", noRewrite)
}
}
func TestResolveAPIKeyModelAliasWithResult_SameBasePreservesSuffix(t *testing.T) {
cfg := &internalconfig.Config{
GeminiKey: []internalconfig.GeminiKey{{
APIKey: "k",
Models: []internalconfig.GeminiModel{{
Name: "gemini-2.5-pro",
Alias: "gemini-2.5-pro(8192)",
ForceMapping: true,
}},
}},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
auth := &Auth{ID: "gemini-auth", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}}
if _, err := mgr.Register(ctx, auth); err != nil {
t.Fatalf("register auth: %v", err)
}
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "gemini-2.5-pro(8192)")
if result.UpstreamModel != "gemini-2.5-pro(8192)" || !result.ForceMapping || result.OriginalAlias != "gemini-2.5-pro(8192)" {
t.Fatalf("resolveAPIKeyModelAliasWithResult() = %+v, want same-base suffix preserved", result)
}
}
func TestResolveAPIKeyModelAliasWithResult_ForceMappingUsesConfigAliasNotRequestSuffix(t *testing.T) {
cfg := &internalconfig.Config{
CodexKey: []internalconfig.CodexKey{{
APIKey: "codex-key",
Models: []internalconfig.CodexModel{{
Name: "gpt-5.5",
Alias: "claude-sonnet-4-5",
ForceMapping: true,
}},
}},
}
mgr := NewManager(nil, nil, nil)
mgr.SetConfig(cfg)
ctx := context.Background()
auth := &Auth{ID: "codex-auth", Provider: "codex", Attributes: map[string]string{"api_key": "codex-key"}}
if _, err := mgr.Register(ctx, auth); err != nil {
t.Fatalf("register auth: %v", err)
}
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "claude-sonnet-4-5(high)")
if result.UpstreamModel != "gpt-5.5(high)" {
t.Fatalf("upstream = %q want gpt-5.5(high)", result.UpstreamModel)
}
if result.OriginalAlias != "claude-sonnet-4-5" {
t.Fatalf("OriginalAlias = %q want claude-sonnet-4-5", result.OriginalAlias)
}
}