| 1 | package config |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "net/http" |
| 7 | "net/http/httptest" |
| 8 | "testing" |
| 9 | ) |
| 10 | |
| 11 | func TestBuildModelFetchURLs(t *testing.T) { |
| 12 | tests := []struct { |
| 13 | name string |
| 14 | base string |
| 15 | override string |
| 16 | want []string |
| 17 | }{ |
| 18 | { |
| 19 | name: "root endpoint keeps legacy models path first", |
| 20 | base: "https://api.deepseek.com", |
| 21 | want: []string{"https://api.deepseek.com/models", "https://api.deepseek.com/v1/models"}, |
| 22 | }, |
| 23 | { |
| 24 | name: "versioned endpoint uses models under version", |
| 25 | base: "https://api.example.com/v1", |
| 26 | want: []string{"https://api.example.com/v1/models"}, |
| 27 | }, |
| 28 | { |
| 29 | name: "non-v1 version keeps v1 fallback", |
| 30 | base: "https://open.bigmodel.cn/api/coding/paas/v4", |
| 31 | want: []string{ |
| 32 | "https://open.bigmodel.cn/api/coding/paas/v4/models", |
| 33 | "https://open.bigmodel.cn/api/coding/paas/v4/v1/models", |
| 34 | }, |
| 35 | }, |
| 36 | { |
| 37 | name: "anthropic compatible subpath adds root candidates", |
| 38 | base: "https://api.deepseek.com/anthropic", |
| 39 | want: []string{ |
| 40 | "https://api.deepseek.com/anthropic/models", |
| 41 | "https://api.deepseek.com/anthropic/v1/models", |
| 42 | "https://api.deepseek.com/models", |
| 43 | "https://api.deepseek.com/v1/models", |
| 44 | }, |
| 45 | }, |
| 46 | { |
| 47 | name: "override wins", |
| 48 | base: "https://api.deepseek.com", |
| 49 | override: "https://api.deepseek.com/custom/models", |
| 50 | want: []string{"https://api.deepseek.com/custom/models"}, |
| 51 | }, |
| 52 | } |
| 53 | for _, tt := range tests { |
| 54 | t.Run(tt.name, func(t *testing.T) { |
| 55 | got, err := BuildModelFetchURLs(tt.base, tt.override) |
| 56 | if err != nil { |
| 57 | t.Fatalf("BuildModelFetchURLs: %v", err) |
| 58 | } |
| 59 | if len(got) != len(tt.want) { |
| 60 | t.Fatalf("got %v, want %v", got, tt.want) |
| 61 | } |
| 62 | for i := range got { |
| 63 | if got[i] != tt.want[i] { |
| 64 | t.Fatalf("got %v, want %v", got, tt.want) |
| 65 | } |
| 66 | } |
| 67 | }) |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | func TestProviderFetchModelsFallsBackToV1Models(t *testing.T) { |
| 72 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 73 | if r.URL.Path == "/models" { |
| 74 | http.NotFound(w, r) |
| 75 | return |
| 76 | } |
| 77 | if r.URL.Path != "/v1/models" { |
| 78 | t.Fatalf("unexpected path %s", r.URL.Path) |
| 79 | } |
| 80 | if r.Header.Get("Authorization") != "Bearer test-key" { |
| 81 | http.Error(w, "bad key", http.StatusUnauthorized) |
| 82 | return |
| 83 | } |
| 84 | _ = json.NewEncoder(w).Encode(map[string]any{ |
| 85 | "data": []map[string]string{{"id": "model-b"}, {"id": "model-a"}}, |
| 86 | }) |
| 87 | })) |
| 88 | defer srv.Close() |
| 89 | |
| 90 | p := ProviderEntry{Name: "test", BaseURL: srv.URL, APIKeyEnv: "FETCH_MODELS_TEST_KEY", resolvedAPIKey: "test-key"} |
| 91 | got, err := p.FetchModels(context.Background()) |
| 92 | if err != nil { |
| 93 | t.Fatalf("FetchModels: %v", err) |
| 94 | } |
| 95 | if len(got) != 2 || got[0] != "model-a" || got[1] != "model-b" { |
| 96 | t.Fatalf("got %v, want [model-a model-b]", got) |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | func TestProviderFetchModelsContinuesAfterRootAuthFailure(t *testing.T) { |
| 101 | var paths []string |
| 102 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 103 | paths = append(paths, r.URL.Path) |
| 104 | switch r.URL.Path { |
| 105 | case "/models": |
| 106 | http.Error(w, `{"error":"wrong endpoint"}`, http.StatusUnauthorized) |
| 107 | case "/v1/models": |
| 108 | if r.Header.Get("Authorization") != "Bearer test-key" { |
| 109 | http.Error(w, "bad key", http.StatusUnauthorized) |
| 110 | return |
| 111 | } |
| 112 | _ = json.NewEncoder(w).Encode(map[string]any{ |
| 113 | "data": []map[string]string{{"id": "model-a"}}, |
| 114 | }) |
| 115 | default: |
| 116 | t.Fatalf("unexpected path %s", r.URL.Path) |
| 117 | } |
| 118 | })) |
| 119 | defer srv.Close() |
| 120 | |
| 121 | p := ProviderEntry{Name: "test", BaseURL: srv.URL, APIKeyEnv: "FETCH_MODELS_TEST_KEY", resolvedAPIKey: "test-key"} |
| 122 | got, err := p.FetchModels(context.Background()) |
| 123 | if err != nil { |
| 124 | t.Fatalf("FetchModels: %v", err) |
| 125 | } |
| 126 | if len(got) != 1 || got[0] != "model-a" { |
| 127 | t.Fatalf("got %v, want [model-a]", got) |
| 128 | } |
| 129 | if len(paths) != 2 || paths[0] != "/models" || paths[1] != "/v1/models" { |
| 130 | t.Fatalf("paths = %v, want [/models /v1/models]", paths) |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | func TestProviderFetchModelsUsesSetupProbeEnv(t *testing.T) { |
| 135 | const key = "FETCH_MODELS_PROBE_KEY" |
| 136 | t.Setenv(key, "probe-key") |
| 137 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 138 | if r.Header.Get("Authorization") != "Bearer probe-key" { |
| 139 | http.Error(w, "bad key", http.StatusUnauthorized) |
| 140 | return |
| 141 | } |
| 142 | _ = json.NewEncoder(w).Encode(map[string]any{ |
| 143 | "data": []map[string]string{{"id": "probe-model"}}, |
| 144 | }) |
| 145 | })) |
| 146 | defer srv.Close() |
| 147 | |
| 148 | p := ProviderEntry{Name: "probe", BaseURL: srv.URL, APIKeyEnv: key} |
| 149 | p.ResolveAPIKeyFromProcessEnvForProbe() |
| 150 | got, err := p.FetchModels(context.Background()) |
| 151 | if err != nil { |
| 152 | t.Fatalf("FetchModels: %v", err) |
| 153 | } |
| 154 | if len(got) != 1 || got[0] != "probe-model" { |
| 155 | t.Fatalf("models = %v, want [probe-model]", got) |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | func TestProviderFetchModelsAllowsNoAuthEndpoint(t *testing.T) { |
| 160 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 161 | if r.Header.Get("Authorization") != "" { |
| 162 | http.Error(w, "unexpected auth header", http.StatusBadRequest) |
| 163 | return |
| 164 | } |
| 165 | _ = json.NewEncoder(w).Encode(map[string]any{ |
| 166 | "data": []map[string]string{{"id": "local-b"}, {"id": "local-a"}}, |
| 167 | }) |
| 168 | })) |
| 169 | defer srv.Close() |
| 170 | |
| 171 | p := ProviderEntry{Name: "local", BaseURL: srv.URL} |
| 172 | got, err := p.FetchModels(context.Background()) |
| 173 | if err != nil { |
| 174 | t.Fatalf("FetchModels no-auth: %v", err) |
| 175 | } |
| 176 | if len(got) != 2 || got[0] != "local-a" || got[1] != "local-b" { |
| 177 | t.Fatalf("got %v, want [local-a local-b]", got) |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | func TestProviderFetchModelsUsesAnthropicAuthMode(t *testing.T) { |
| 182 | tests := []struct { |
| 183 | name string |
| 184 | authHeader bool |
| 185 | assertAuth func(t *testing.T, r *http.Request) |
| 186 | }{ |
| 187 | { |
| 188 | name: "x-api-key", |
| 189 | authHeader: false, |
| 190 | assertAuth: func(t *testing.T, r *http.Request) { |
| 191 | t.Helper() |
| 192 | if got := r.Header.Get("x-api-key"); got != "anthropic-key" { |
| 193 | t.Fatalf("x-api-key = %q, want anthropic-key", got) |
| 194 | } |
| 195 | if got := r.Header.Get("Authorization"); got != "" { |
| 196 | t.Fatalf("Authorization = %q, want omitted", got) |
| 197 | } |
| 198 | }, |
| 199 | }, |
| 200 | { |
| 201 | name: "bearer", |
| 202 | authHeader: true, |
| 203 | assertAuth: func(t *testing.T, r *http.Request) { |
| 204 | t.Helper() |
| 205 | if got := r.Header.Get("Authorization"); got != "Bearer anthropic-key" { |
| 206 | t.Fatalf("Authorization = %q, want Bearer anthropic-key", got) |
| 207 | } |
| 208 | if got := r.Header.Get("x-api-key"); got != "" { |
| 209 | t.Fatalf("x-api-key = %q, want omitted", got) |
| 210 | } |
| 211 | }, |
| 212 | }, |
| 213 | } |
| 214 | |
| 215 | for _, tt := range tests { |
| 216 | t.Run(tt.name, func(t *testing.T) { |
| 217 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 218 | if r.URL.Path != "/anthropic/models" { |
| 219 | t.Fatalf("unexpected path %s", r.URL.Path) |
| 220 | } |
| 221 | tt.assertAuth(t, r) |
| 222 | _ = json.NewEncoder(w).Encode(map[string]any{ |
| 223 | "data": []map[string]string{{"id": "anthropic-model"}}, |
| 224 | }) |
| 225 | })) |
| 226 | defer srv.Close() |
| 227 | |
| 228 | p := ProviderEntry{ |
| 229 | Name: "anthropic-compatible", |
| 230 | Kind: "anthropic", |
| 231 | BaseURL: srv.URL + "/anthropic", |
| 232 | APIKeyEnv: "ANTHROPIC_COMPATIBLE_KEY", |
| 233 | AuthHeader: tt.authHeader, |
| 234 | resolvedAPIKey: "anthropic-key", |
| 235 | } |
| 236 | got, err := p.FetchModels(context.Background()) |
| 237 | if err != nil { |
| 238 | t.Fatalf("FetchModels: %v", err) |
| 239 | } |
| 240 | if len(got) != 1 || got[0] != "anthropic-model" { |
| 241 | t.Fatalf("got %v, want [anthropic-model]", got) |
| 242 | } |
| 243 | }) |
| 244 | } |
| 245 | } |
| 246 |