ModelRouter/internal/provider/model_kind_test.go

package provider

import (
	"encoding/json"
	"testing"

	"llmsproxy/internal/config"
)

// TestSourceStatus_CarriesModelKind is the fix for a UI bug where the chat
// test page listed image models next to chat ones.
//
// The root cause was structural, not a missing filter: SourceStatus.Models was
// a flat []string, so by the time the UI had the list there was nothing left to
// tell "gpt-image-2" from "claude-sonnet-5.5". Picking the former sent the
// request to /api/chat, which cannot serve it, and the failure said nothing
// about the real reason.
//
// Models stays a []string for callers that only need names; ModelKinds is the
// additive field that carries the missing dimension.
func TestSourceStatus_CarriesModelKind(t *testing.T) {
	src := config.Source{
		Name:     "vela",
		BaseURL:  "https://example.invalid",
		Adapter:  "openai",
		APIKey:   "sk-test",
		Endpoint: "/chat/completions",
		Models: []config.Model{
			{ID: "claude-sonnet-5.5"},          // kind omitted -> chat default
			{ID: "gpt-image-2", Kind: "image"}, // image
			{ID: "seedream-5.0", Kind: "image"},
		},
	}
	p := newTestProvider(t, src)

	kinds := p.ModelsWithKind()
	if len(kinds) != 3 {
		t.Fatalf("want 3 entries, got %d", len(kinds))
	}

	// Order must mirror Models() so the UI can zip the two lists.
	names := p.Models()
	for i := range kinds {
		if kinds[i].ID != names[i] {
			t.Errorf("index %d: ModelKinds has %q but Models has %q — the lists must stay aligned",
				i, kinds[i].ID, names[i])
		}
	}

	want := map[string]string{
		"claude-sonnet-5.5": "chat",
		"gpt-image-2":       "image",
		"seedream-5.0":      "image",
	}
	for _, k := range kinds {
		if got := want[k.ID]; k.Kind != got {
			t.Errorf("%s: kind = %q, want %q", k.ID, k.Kind, got)
		}
	}
}

// TestModelKinds_CountsImageModelsInARealCatalog: the number of image models is
// what decides whether the UI needs the split at all, so assert against a
// catalog with both kinds rather than a hand-picked pair.
func TestModelKinds_CountsImageModelsInARealCatalog(t *testing.T) {
	src := config.Source{
		Name: "vela", BaseURL: "https://example.invalid", Adapter: "openai", APIKey: "k",
		Models: []config.Model{
			{ID: "a"}, {ID: "b"}, {ID: "c"}, {ID: "d"},
			{ID: "img1", Kind: "image"}, {ID: "img2", Kind: "image"},
		},
	}
	p := newTestProvider(t, src)
	var images int
	for _, k := range p.ModelsWithKind() {
		if k.Kind == "image" {
			images++
		}
	}
	if images != 2 {
		t.Errorf("image models = %d, want 2", images)
	}
}

// TestSourceStatusJSON_ModelKindsOmittedWhenEmpty pins the backward-compat
// promise: an older UI parsing this JSON must not break on a new field, and a
// source with no model list must not gain a "model_kinds": null.
func TestSourceStatusJSON_ModelKindsOmittedWhenEmpty(t *testing.T) {
	b, err := json.Marshal(SourceStatus{Name: "x", Models: []string{"a"}})
	if err != nil {
		t.Fatal(err)
	}
	var m map[string]interface{}
	if err := json.Unmarshal(b, &m); err != nil {
		t.Fatal(err)
	}
	if _, present := m["model_kinds"]; present {
		t.Errorf("model_kinds present when empty: %s — should be omitted for older UIs", b)
	}
	if _, present := m["models"]; !present {
		t.Errorf("models field disappeared: %s", b)
	}
}

// newTestProvider (defined in provider_test.go) spins up a real Lua VM, which
// is heavier than these tests need, but reusing it keeps one construction path
// for the whole package.