291 lines
9.4 KiB
Go
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)
|
|
}
|
|
}
|