ModelRouter/internal/core/core_test.go

package core

import (
	"path/filepath"
	"sync"
	"testing"

	"llmsproxy/internal/config"
)

// newTestConfig returns an in-memory config backed by temp files so
// NewFromConfig (store load, key/auto seeding, registry build) is exercised
// like production.
func newTestConfig(t *testing.T) *config.Config {
	t.Helper()
	dir := t.TempDir()
	return &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-gw-test"},
		Sources: []config.Source{
			{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}}},
		},
	}
}

func newTestCore(t *testing.T, cfg *config.Config) *Core {
	t.Helper()
	c, err := NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	return c
}

func TestCleanScopesNilForEmpty(t *testing.T) {
	if got := cleanScopes(nil); got != nil {
		t.Fatalf("nil in must yield nil out, got %#v", got)
	}
	if got := cleanScopes([]config.ModelScope{{Model: ""}}); got != nil {
		t.Fatalf("all-empty entries must yield nil, got %#v", got)
	}
	got := cleanScopes([]config.ModelScope{{Model: "m", Source: "undefined"}})
	if len(got) != 1 || got[0].Source != "" || got[0].Model != "m" {
		t.Fatalf("placeholder normalization wrong: %#v", got)
	}
}

// TestCreateKeyWithoutScopeUnrestricted: a user key created without a model
// scope must be unrestricted in-process (nil), identical to what config.yaml
// omitempty produces after a restart — no 403 lock-out, no silent widening.
func TestCreateKeyWithoutScopeUnrestricted(t *testing.T) {
	c := newTestCore(t, newTestConfig(t))
	k, err := c.CreateKey("svc", "user", nil, "")
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	rec, ok := c.FindKey(k.Key)
	if !ok {
		t.Fatal("key not found")
	}
	if rec.Models != nil {
		t.Fatalf("user key without scope must be unrestricted, got %#v", rec.Models)
	}
}

// TestUpdateKeyPreservesScopeWhenOmitted: PUT /api/keys without a models field
// must keep the existing scope; only an explicit [] clears it.
func TestUpdateKeyPreservesScopeWhenOmitted(t *testing.T) {
	c := newTestCore(t, newTestConfig(t))
	k, err := c.CreateKey("svc", "user", []config.ModelScope{{Model: "gpt-4o"}}, "")
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	upd, err := c.UpdateKey(k.Key, "renamed", "", nil, "")
	if err != nil {
		t.Fatalf("update: %v", err)
	}
	if len(upd.Models) != 1 || upd.Models[0].Model != "gpt-4o" {
		t.Fatalf("scope must survive an update without models, got %#v", upd.Models)
	}
	cleared, err := c.UpdateKey(k.Key, "", "", []config.ModelScope{}, "")
	if err != nil {
		t.Fatalf("clear: %v", err)
	}
	if cleared.Models != nil {
		t.Fatalf("explicit empty list must clear the scope, got %#v", cleared.Models)
	}
}

// TestSaveAutoRulesNormalizesModelCase: a rule whose model name differs only
// in case must be normalized to the exact configured spelling so cooldown,
// quota windows and the upstream model id all agree.
func TestSaveAutoRulesNormalizesModelCase(t *testing.T) {
	c := newTestCore(t, newTestConfig(t))
	if err := c.SaveAutoRules([]config.ModelScope{{Model: "GPT-4O", Source: "s1", Tier: 1}}); err != nil {
		t.Fatalf("save: %v", err)
	}
	ch := c.AutoChain()
	if ch == nil || len(ch.Tiers) == 0 || len(ch.Tiers[0].Slots) == 0 {
		t.Fatalf("chain empty: %#v", ch)
	}
	if got := ch.Tiers[0].Slots[0].Model; got != "gpt-4o" {
		t.Fatalf("slot model must be normalized to the config spelling, got %q", got)
	}
	if got := ch.Tiers[0].Slots[0].Source; got != "s1" {
		t.Fatalf("slot source wrong: %q", got)
	}
}

// TestConcurrentKeyMutation exercises the c.cfg.Keys lock under -race: key
// creation/lookup/deletion from many goroutines at once.
func TestConcurrentKeyMutation(t *testing.T) {
	c := newTestCore(t, newTestConfig(t))
	var wg sync.WaitGroup
	for i := 0; i < 16; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			k, err := c.CreateKey("k", "user", nil, "")
			if err == nil {
				if rec, ok := c.FindKey(k.Key); ok && rec.Key != k.Key {
					t.Errorf("wrong key returned")
				}
				c.UpdateKey(k.Key, "renamed", "", nil, "")
				c.DeleteKey(k.Key)
			}
		}()
	}
	wg.Wait()
}