ModelRouter/internal/gateway/discover_api_test.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"llmsproxy/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
	}
}