localmodelroutex/backend/internal/catalog/cache_test.go

package catalog

import (
	"testing"

	"github.com/google/uuid"
)

func newTestCache(strategy Strategy, entries ...*Entry) *Cache {
	c := &Cache{byAddr: map[string]*Entry{}, byID: map[uuid.UUID]*Entry{}, standby: map[string]*Entry{}, strategy: strategy}
	for _, e := range entries {
		if e.Address == "" {
			e.Address = e.ProviderName + "/" + e.Slug
		}
		if e.ID == uuid.Nil {
			e.ID = uuid.New()
		}
		c.byAddr[e.Address] = e
		c.byID[e.ID] = e
		c.entries = append(c.entries, e)
	}
	return c
}

// addStandby registers an enabled model the provider does not serve right now.
func addStandby(c *Cache, e *Entry) *Entry {
	if e.Address == "" {
		e.Address = e.ProviderName + "/" + e.Slug
	}
	if e.ID == uuid.Nil {
		e.ID = uuid.New()
	}
	c.standby[e.Address] = e
	return e
}

func backup(e *Entry) *uuid.UUID { id := e.ID; return &id }

func addresses(entries []*Entry) []string {
	out := make([]string, len(entries))
	for i, e := range entries {
		out[i] = e.Address
	}
	return out
}

func TestRouteFollowsBackupChainAndSkipsDownBackups(t *testing.T) {
	a := &Entry{ProviderName: "local", Slug: "a"}
	b := &Entry{ProviderName: "local", Slug: "b"}
	cEntry := &Entry{ProviderName: "cloud", Slug: "c"}
	c := newTestCache(StrategyRoundRobin, a, cEntry)
	addStandby(c, b)
	a.FallbackID = backup(b)
	b.FallbackID = backup(cEntry)

	r, err := c.Route("local/a")
	if err != nil {
		t.Fatal(err)
	}
	if got := addresses(r.Targets); len(got) != 2 || got[0] != "local/a" || got[1] != "cloud/c" {
		t.Fatalf("targets = %v, want [local/a cloud/c]", got)
	}
	if r.Primary != "local/a" || r.Slug != "a" || r.IsFallback(a) || !r.IsFallback(cEntry) {
		t.Errorf("route = %+v", r)
	}
}

func TestRouteUnavailableModelGoesToBackup(t *testing.T) {
	b := &Entry{ProviderName: "cloud", Slug: "org/a"}
	c := newTestCache(StrategyRoundRobin, b)
	a := addStandby(c, &Entry{ProviderName: "local", Slug: "a"})
	a.FallbackID = backup(b)

	r, err := c.Route("local/a")
	if err != nil {
		t.Fatal(err)
	}
	if got := addresses(r.Targets); len(got) != 1 || got[0] != "cloud/org/a" {
		t.Fatalf("targets = %v", got)
	}
	// The client asked for local/a, so the allowlist and the reported model
	// use it, not the backup.
	if r.Address != "local/a" || r.Primary != "local/a" || r.Slug != "a" || !r.IsFallback(b) {
		t.Errorf("route = %+v", r)
	}
	if e, err := c.Resolve("local/a"); err != nil || e != b {
		t.Errorf("Resolve = %v, %v", e, err)
	}
	if got := addresses(c.List()); len(got) != 2 {
		t.Errorf("List must include the down model with a live backup: %v", got)
	}
}

func TestRouteUnavailableModelWithoutBackupIsNotFound(t *testing.T) {
	c := newTestCache(StrategyRoundRobin)
	addStandby(c, &Entry{ProviderName: "local", Slug: "a"})
	if _, err := c.Route("local/a"); err == nil {
		t.Fatal("expected not found")
	}
	if got := c.List(); len(got) != 0 {
		t.Errorf("List = %v", addresses(got))
	}
}

func TestRouteMutualBackupsDoNotLoop(t *testing.T) {
	a := &Entry{ProviderName: "x", Slug: "a"}
	b := &Entry{ProviderName: "x", Slug: "b"}
	c := newTestCache(StrategyRoundRobin, a, b)
	a.FallbackID = backup(b)
	b.FallbackID = backup(a)
	r, err := c.Route("x/b")
	if err != nil {
		t.Fatal(err)
	}
	if got := addresses(r.Targets); len(got) != 2 || got[0] != "x/b" || got[1] != "x/a" {
		t.Fatalf("targets = %v", got)
	}
}

func TestRouteAnyCarriesBackups(t *testing.T) {
	a := &Entry{ProviderName: "openai", Slug: "gpt-4o"}
	b := &Entry{ProviderName: "cloud", Slug: "other"}
	c := newTestCache(StrategyRoundRobin, a, b)
	a.FallbackID = backup(b)
	r, err := c.Route("any/gpt-4o")
	if err != nil {
		t.Fatal(err)
	}
	if got := addresses(r.Targets); len(got) != 2 || got[1] != "cloud/other" {
		t.Fatalf("targets = %v", got)
	}
	if r.Address != "any/gpt-4o" || r.Primary != "openai/gpt-4o" {
		t.Errorf("route = %+v", r)
	}
}

func TestSplitAddress(t *testing.T) {
	cases := []struct {
		in             string
		provider, slug string
		ok             bool
	}{
		{"openai/gpt-4o", "openai", "gpt-4o", true},
		{"openrouter/openai/gpt-4o", "openrouter", "openai/gpt-4o", true},
		{"any/gpt-4o", "any", "gpt-4o", true},
		{"noslash", "", "", false},
		{"/leading", "", "", false},
		{"trailing/", "", "", false},
	}
	for _, c := range cases {
		p, s, ok := SplitAddress(c.in)
		if ok != c.ok || p != c.provider || s != c.slug {
			t.Errorf("SplitAddress(%q) = (%q,%q,%v), want (%q,%q,%v)", c.in, p, s, ok, c.provider, c.slug, c.ok)
		}
	}
}

func TestResolveDirect(t *testing.T) {
	c := newTestCache(StrategyRoundRobin, &Entry{ProviderName: "openai", Slug: "gpt-4o"})
	e, err := c.Resolve("openai/gpt-4o")
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if e.ProviderName != "openai" || e.Slug != "gpt-4o" {
		t.Fatalf("got %+v", e)
	}
	if _, err := c.Resolve("openai/missing"); err == nil {
		t.Fatal("expected error for unknown model")
	}
}

func TestResolveAnyRoundRobin(t *testing.T) {
	// Two providers serve gpt-4o (one with the slug provider-prefixed).
	c := newTestCache(StrategyRoundRobin,
		&Entry{ProviderName: "openai", Slug: "gpt-4o"},
		&Entry{ProviderName: "openrouter", Slug: "openai/gpt-4o"},
	)
	seen := map[string]int{}
	for i := 0; i < 4; i++ {
		e, err := c.Resolve("any/gpt-4o")
		if err != nil {
			t.Fatalf("Resolve: %v", err)
		}
		seen[e.ProviderName]++
	}
	if seen["openai"] != 2 || seen["openrouter"] != 2 {
		t.Fatalf("round-robin distribution = %v, want 2/2", seen)
	}
}

func TestResolveAnyLowestCost(t *testing.T) {
	c := newTestCache(StrategyLowestCost,
		&Entry{ProviderName: "pricey", Slug: "gpt-4o", InputPricePerMTok: 10, OutputPricePerMTok: 30},
		&Entry{ProviderName: "cheap", Slug: "gpt-4o", InputPricePerMTok: 2, OutputPricePerMTok: 6},
	)
	for i := 0; i < 3; i++ {
		e, err := c.Resolve("any/gpt-4o")
		if err != nil {
			t.Fatalf("Resolve: %v", err)
		}
		if e.ProviderName != "cheap" {
			t.Fatalf("lowest_cost picked %q, want cheap", e.ProviderName)
		}
	}
}

func TestResolveAnyNoCandidates(t *testing.T) {
	c := newTestCache(StrategyRoundRobin, &Entry{ProviderName: "openai", Slug: "gpt-4o"})
	if _, err := c.Resolve("any/llama-3"); err == nil {
		t.Fatal("expected error when no provider serves the model")
	}
}