package gateway import ( "context" "encoding/json" "errors" "net/http" "net/http/httptest" "strings" "testing" "modelrouter/internal/config" ) // upstream that answers /v1/models with the real captured catalog func catalogUpstream(t *testing.T) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" { w.WriteHeader(http.StatusNotFound) return } if got := r.Header.Get("Authorization"); got != "Bearer sk-secret" { w.WriteHeader(http.StatusUnauthorized) return } _, _ = w.Write([]byte(realCatalogBody)) })) t.Cleanup(srv.Close) return srv } func TestDefaultSourceName(t *testing.T) { cases := map[string]string{ "https://api.vendor.com/v1": "vendor.com", "https://api.openai.com/v1": "openai.com", "https://www.vendor.com/v1": "vendor.com", "http://127.0.0.1:8080/v1": "127.0.0.1", "https://amr-link.open-design.ai/v1": "amr-link.open-design.ai", "not a url at all": "", } for in, want := range cases { if got := defaultSourceName(in); got != want { t.Errorf("defaultSourceName(%q) = %q, want %q", in, got, want) } } } func TestSelectModels_IncludeUnknownIsAnError(t *testing.T) { chat := []string{"a", "b"} image := []string{"img"} // Silently dropping a typo'd id would create a source that looks // configured and 404s on every request — much worse than a rejection. _, _, _, err := selectModels(chat, image, []string{"a", "typo"}, nil) if err == nil { t.Fatal("expected an error for an include id that is not in the catalog") } if !strings.Contains(err.Error(), "typo") { t.Errorf("error should name the offending id, got %q", err) } } func TestSelectModels_IncludeNarrows(t *testing.T) { chat, image, _, err := selectModels([]string{"a", "b", "c"}, []string{"img1", "img2"}, []string{"b", "img1"}, nil) if err != nil { t.Fatal(err) } if len(chat) != 1 || chat[0] != "b" { t.Errorf("chat = %v, want [b]", chat) } if len(image) != 1 || image[0] != "img1" { t.Errorf("image = %v, want [img1]", image) } } func TestSelectModels_ExcludeUnknownIsFine(t *testing.T) { chat, _, _, err := selectModels([]string{"a"}, nil, nil, []string{"never-existed"}) if err != nil { t.Errorf("excluding a model that is not present should be allowed: %v", err) } if len(chat) != 1 { t.Errorf("chat = %v, want [a] unchanged", chat) } } func TestSelectModels_NeverReturnsNil(t *testing.T) { chat, image, warns, err := selectModels(nil, nil, nil, nil) if err != nil { t.Fatal(err) } if chat == nil || image == nil { t.Fatal("nil slices would serialize as JSON null and break the UI") } if len(warns) == 0 { t.Error("an empty catalog deserves a warning, not silence") } } func TestSelectModels_WarnsAboutImageOnlyCatalog(t *testing.T) { _, _, warns, err := selectModels(nil, []string{"gpt-image-2"}, nil, nil) if err != nil { t.Fatal(err) } joined := strings.Join(warns, " ") if !strings.Contains(joined, "AUTO") { t.Errorf("warnings %v should say AUTO has nothing to select", warns) } } func TestSelectModels_WarnsThatImagesAreNotInAUTO(t *testing.T) { _, _, warns, err := selectModels([]string{"a"}, []string{"img"}, nil, nil) if err != nil { t.Fatal(err) } if !strings.Contains(strings.Join(warns, " "), "images/generations") { t.Errorf("warnings %v should point at the image endpoint", warns) } } // TestDiscoveryResultAndConfigRoundTrip checks the fields the WebUI will echo // back actually survive into a config.Source — a mismatch here means the // operator previews one set of values and saves another. func TestDiscoveryResultAndConfigRoundTrip(t *testing.T) { res := &discoveryResult{ BaseURL: "https://h", Endpoint: "/v1/chat/completions", Models: []string{"m1", "m2"}, ImageModels: []string{"img1"}, } src := config.Source{ Name: "x", BaseURL: res.BaseURL, Endpoint: res.Endpoint, Models: modelsToConfig(res.Models, res.ImageModels), } if len(src.Models) != 3 { t.Fatalf("models = %d, want 3", len(src.Models)) } if src.Endpoint != "/v1/chat/completions" { t.Errorf("endpoint lost: %q", src.Endpoint) } var img int for _, m := range src.Models { if m.Kind == "image" { img++ } } if img != 1 { t.Errorf("image models = %d, want 1", img) } } // TestDiscoverModels_SendsAuthorizationHeader pins the only way the feature // works against a real vendor: without the bearer token every paid /models // endpoint returns 401. func TestDiscoverModels_SendsAuthorizationHeader(t *testing.T) { var seen string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { seen = r.Header.Get("Authorization") _, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`)) })) defer srv.Close() if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-secret"); err != nil { t.Fatal(err) } if seen != "Bearer sk-secret" { t.Errorf("Authorization = %q, want %q", seen, "Bearer sk-secret") } } // TestDiscoverModels_WrongKeyIsUnauthorizedNotBadGateway pins the distinction // that decides whether the operator knows what to do next. // // A rejected key and an unreachable vendor are the two most common discovery // failures, and they need opposite actions: re-type the key, versus go check the // vendor's status page. Both used to surface as one generic error, so the // operator had to guess. func TestDiscoverModels_WrongKeyIsUnauthorizedNotBadGateway(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusUnauthorized) _, _ = w.Write([]byte(`{"error":{"message":"invalid api key","code":"invalid_api_key"}}`)) })) defer srv.Close() _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-wrong") if err == nil { t.Fatal("expected a failure against a 401 endpoint") } var f *discoveryFailure if !errors.As(err, &f) { t.Fatalf("error should be classified, got %T: %v", err, err) } if got := f.StatusOrDefault(); got != http.StatusUnauthorized { t.Errorf("status = %d, want %d — a wrong key must not read as an upstream outage", got, http.StatusUnauthorized) } // The message has to name which credential is at fault, or the status code // alone still sends the operator to the wrong place. if !strings.Contains(strings.ToLower(f.Error()), "api key") { t.Errorf("message %q should name the api key as the cause", f.Error()) } if !strings.Contains(f.Error(), "/v1/models") { t.Errorf("message %q should list what was tried", f.Error()) } } // TestDiscoverModels_UnreachableStaysBadGateway: the classification must not // over-correct — an unresolvable host is still the vendor's or the network's // problem, not the key's. func TestDiscoverModels_UnreachableStaysBadGateway(t *testing.T) { _, err := DiscoverModels(context.Background(), "https://no-such-host-xyz.invalid/v1", "k") if err == nil { t.Fatal("expected a failure for an unresolvable host") } var f *discoveryFailure if !errors.As(err, &f) { t.Fatalf("error should be classified, got %T", err) } if got := f.StatusOrDefault(); got != http.StatusBadGateway { t.Errorf("status = %d, want %d", got, http.StatusBadGateway) } if !strings.Contains(strings.ToLower(f.Error()), "resolve") { t.Errorf("message %q should say the hostname does not resolve", f.Error()) } } // TestDiscoveryCandidates_NoDuplicateRequests: base "https://h/v1" with path // "/models" and base "https://h" with path "/v1/models" request the SAME URL. // Listing it twice in a failure message reads as "two things were tried" when // only one was. func TestDiscoveryCandidates_NoDuplicateRequests(t *testing.T) { seen := map[string]bool{} for _, c := range discoveryCandidates("https://h/v1") { url := c.baseURL + c.modelsPath if seen[url] { t.Errorf("candidate %q would be requested twice", url) } seen[url] = true } if len(seen) < 2 { t.Errorf("only %d distinct URLs; the search collapsed", len(seen)) } } // TestDiscoverModels_FailureListsEachURLOnce is the end-to-end form of the same // rule, on the error text an operator actually reads. func TestDiscoverModels_FailureListsEachURLOnce(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k") if err == nil { t.Fatal("expected failure") } msg := err.Error() needle := srv.URL + "/v1/models" if n := strings.Count(msg, needle); n > 1 { t.Errorf("URL %s appears %d times in %q", needle, n, msg) } } func TestDiscoverModels_EmptyKeyOmitsHeader(t *testing.T) { var present bool srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, present = r.Header["Authorization"] _, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`)) })) defer srv.Close() if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", ""); err != nil { t.Fatal(err) } if present { t.Error("an empty key should omit the header, not send an empty bearer") } } // TestDiscoverModels_RealCatalogEndToEnd walks the actual captured response // through the whole path and asserts the numbers an operator would see. func TestDiscoverModels_RealCatalogEndToEnd(t *testing.T) { srv := catalogUpstream(t) res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-secret") if err != nil { t.Fatalf("discovery failed: %v", err) } if res.BaseURL != srv.URL+"/v1" { t.Errorf("base = %q, want the typed URL unchanged", res.BaseURL) } if res.Endpoint != "" { t.Errorf("endpoint = %q, want empty (base already carries /v1)", res.Endpoint) } if len(res.Models) != 2 || len(res.ImageModels) != 2 { t.Errorf("chat/image = %d/%d, want 2/2", len(res.Models), len(res.ImageModels)) } if res.Adapter != "openai" { t.Errorf("adapter = %q, want openai", res.Adapter) } // The response must be JSON the UI can consume directly. b, _ := json.Marshal(res) var round map[string]interface{} if err := json.Unmarshal(b, &round); err != nil { t.Fatal(err) } if round["ok"] != nil { // discoverResponse adds ok; discoveryResult does not. Both marshal // fine, so this is only a shape note. _ = round } }