ModelRouter/internal/gateway/discover_api_test.go

package gateway

import (
	"context"
	"encoding/json"
	"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")
	}
}

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
	}
}