package auth import ( "testing" internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config" ) func TestResolveOAuthUpstreamModel_SuffixPreservation(t *testing.T) { t.Parallel() tests := []struct { name string mappings map[string][]internalconfig.ModelNameMapping channel string input string want string }{ { name: "numeric suffix preserved", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro(8192)", want: "gemini-2.5-pro-exp-03-25(8192)", }, { name: "level suffix preserved", mappings: map[string][]internalconfig.ModelNameMapping{ "claude": {{Name: "claude-sonnet-4-5-20250514", Alias: "claude-sonnet-4-5"}}, }, channel: "claude", input: "claude-sonnet-4-5(high)", want: "claude-sonnet-4-5-20250514(high)", }, { name: "no suffix unchanged", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro", want: "gemini-2.5-pro-exp-03-25", }, { name: "config suffix takes priority", mappings: map[string][]internalconfig.ModelNameMapping{ "claude": {{Name: "claude-sonnet-4-5-20250514(low)", Alias: "claude-sonnet-4-5"}}, }, channel: "claude", input: "claude-sonnet-4-5(high)", want: "claude-sonnet-4-5-20250514(low)", }, { name: "auto suffix preserved", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro(auto)", want: "gemini-2.5-pro-exp-03-25(auto)", }, { name: "none suffix preserved", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro(none)", want: "gemini-2.5-pro-exp-03-25(none)", }, { name: "case insensitive alias lookup with suffix", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "Gemini-2.5-Pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro(high)", want: "gemini-2.5-pro-exp-03-25(high)", }, { name: "no mapping returns empty", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "unknown-model(high)", want: "", }, { name: "wrong channel returns empty", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "claude", input: "gemini-2.5-pro(high)", want: "", }, { name: "empty suffix filtered out", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro()", want: "gemini-2.5-pro-exp-03-25", }, { name: "incomplete suffix treated as no suffix", mappings: map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro(high"}}, }, channel: "gemini-cli", input: "gemini-2.5-pro(high", want: "gemini-2.5-pro-exp-03-25", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() mgr := NewManager(nil, nil, nil) mgr.SetConfig(&internalconfig.Config{}) mgr.SetOAuthModelMappings(tt.mappings) auth := createAuthForChannel(tt.channel) got := mgr.resolveOAuthUpstreamModel(auth, tt.input) if got != tt.want { t.Errorf("resolveOAuthUpstreamModel(%q) = %q, want %q", tt.input, got, tt.want) } }) } } func createAuthForChannel(channel string) *Auth { switch channel { case "gemini-cli": return &Auth{Provider: "gemini-cli"} case "claude": return &Auth{Provider: "claude", Attributes: map[string]string{"auth_kind": "oauth"}} case "vertex": return &Auth{Provider: "vertex", Attributes: map[string]string{"auth_kind": "oauth"}} case "codex": return &Auth{Provider: "codex", Attributes: map[string]string{"auth_kind": "oauth"}} case "aistudio": return &Auth{Provider: "aistudio"} case "antigravity": return &Auth{Provider: "antigravity"} case "qwen": return &Auth{Provider: "qwen"} case "iflow": return &Auth{Provider: "iflow"} default: return &Auth{Provider: channel} } } func TestApplyOAuthModelMapping_SuffixPreservation(t *testing.T) { t.Parallel() mappings := map[string][]internalconfig.ModelNameMapping{ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}}, } mgr := NewManager(nil, nil, nil) mgr.SetConfig(&internalconfig.Config{}) mgr.SetOAuthModelMappings(mappings) auth := &Auth{ID: "test-auth-id", Provider: "gemini-cli"} metadata := map[string]any{"existing": "value"} resolvedModel, resultMeta := mgr.applyOAuthModelMapping(auth, "gemini-2.5-pro(8192)", metadata) if resolvedModel != "gemini-2.5-pro-exp-03-25(8192)" { t.Errorf("applyOAuthModelMapping() model = %q, want %q", resolvedModel, "gemini-2.5-pro-exp-03-25(8192)") } originalModel, ok := resultMeta["model_mapping_original_model"].(string) if !ok || originalModel != "gemini-2.5-pro(8192)" { t.Errorf("applyOAuthModelMapping() metadata[model_mapping_original_model] = %v, want %q", resultMeta["model_mapping_original_model"], "gemini-2.5-pro(8192)") } if resultMeta["existing"] != "value" { t.Errorf("applyOAuthModelMapping() metadata[existing] = %v, want %q", resultMeta["existing"], "value") } }