ModelRouter/internal/config/config_test.go

package config

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

func TestLoadAndApplyDefaults(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cfg.yaml")
	content := `
listen: 127.0.0.1:9999
gateway_keys: [sk-1]
default_model: AUTO
adapter_dir: adapters
runtime_file: runtime.json
sources:
  - name: deepseek
    base_url: https://api.deepseek.com
    api_key: sk-d
    adapter: deepseek
    models:
      - id: deepseek-v4-flash
        priority: 100
  - name: ollama
    base_url: http://127.0.0.1:11434
    adapter: ollama
    endpoint: /api/chat
    models:
      - id: llama3
        priority: 50
`
	if err := os.WriteFile(path, []byte(content), 0644); err != nil {
		t.Fatal(err)
	}
	cfg, err := Load(path)
	if err != nil {
		t.Fatalf("load: %v", err)
	}
	if len(cfg.Sources) != 2 {
		t.Fatalf("sources = %d", len(cfg.Sources))
	}
	if cfg.Sources[1].Endpoint != "/api/chat" {
		t.Fatalf("endpoint = %q", cfg.Sources[1].Endpoint)
	}
	if cfg.Sources[1].Timeout == 0 {
		t.Fatal("default timeout not applied")
	}
	if cfg.Sources[1].MaxConcurrent == 0 {
		t.Fatal("default max_concurrent not applied")
	}
	if cfg.Sources[0].Models[0].Priority != 100 {
		t.Fatalf("priority = %d", cfg.Sources[0].Models[0].Priority)
	}
	if cfg.DefaultModel != "AUTO" {
		t.Fatalf("default model = %q", cfg.DefaultModel)
	}
}

func TestEnsureDefaultGeneratesOnMissingFile(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "nested", "config.yaml")
	created, err := EnsureDefault(path)
	if err != nil {
		t.Fatalf("ensure: %v", err)
	}
	if !created {
		t.Fatal("expected creation for missing file")
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read generated: %v", err)
	}
	if fi, err := os.Stat(path); err != nil || (runtime.GOOS != "windows" && fi.Mode().Perm() != 0600) {
		t.Fatalf("generated config must be 0600 (holds plaintext key), got %v", fi.Mode().Perm())
	}
	// generated file must load and be safe-by-default
	cfg, err := Load(path)
	if err != nil {
		t.Fatalf("load generated: %v", err)
	}
	if len(cfg.GatewayKeys) != 1 || !strings.HasPrefix(cfg.GatewayKeys[0], "sk-gw-") {
		t.Fatalf("gateway keys = %v", cfg.GatewayKeys)
	}
	if strings.Contains(string(raw), cfg.GatewayKeys[0]) == false {
		t.Fatal("generated key must be written into the file")
	}
	if cfg.Listen != "127.0.0.1:8080" {
		t.Fatalf("listen = %q, want loopback-only", cfg.Listen)
	}
	if len(cfg.Sources) != 0 {
		t.Fatalf("default config must ship no sources (user adds them), got %+v", cfg.Sources)
	}
	// adapter_dir / runtime_file resolve next to the config file
	if !strings.HasPrefix(cfg.AdapterDir, dir) || !strings.HasPrefix(cfg.RuntimeFile, dir) {
		t.Fatalf("paths must live next to the config: adapter_dir=%s runtime_file=%s", cfg.AdapterDir, cfg.RuntimeFile)
	}
}

func TestEnsureDefaultDoesNotOverwriteExisting(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cfg.yaml")
	content := "listen: 127.0.0.1:9999\ngateway_keys: [sk-keep]\n"
	if err := os.WriteFile(path, []byte(content), 0644); err != nil {
		t.Fatal(err)
	}
	created, err := EnsureDefault(path)
	if err != nil {
		t.Fatalf("ensure: %v", err)
	}
	if created {
		t.Fatal("existing file must not be reported as created")
	}
	raw, _ := os.ReadFile(path)
	if string(raw) != content {
		t.Fatalf("existing file was overwritten: %q", raw)
	}
}

func TestNewGatewayKeyIsRandom(t *testing.T) {
	a, err := NewGatewayKey()
	if err != nil {
		t.Fatal(err)
	}
	b, err := NewGatewayKey()
	if err != nil {
		t.Fatal(err)
	}
	if a == b {
		t.Fatal("two generated keys must differ")
	}
	if !strings.HasPrefix(a, "sk-gw-") || len(a) != len("sk-gw-")+32 {
		t.Fatalf("unexpected key format: %q", a)
	}
}

func TestApplyDefaultsDuplicateSource(t *testing.T) {
	cfg := Config{Sources: []Source{
		{Name: "a", BaseURL: "http://x", Models: []Model{{ID: "m1"}}},
		{Name: "a", BaseURL: "http://y", Models: []Model{{ID: "m2"}}},
	}}
	if err := cfg.ApplyDefaults(); err == nil {
		t.Fatal("expected duplicate source error")
	}
}

func TestApplyDefaultsNoSources(t *testing.T) {
	// empty source list is valid (sources may be added later via the Web UI)
	cfg := Config{}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	if cfg.Listen != ":8080" || cfg.DefaultModel != "AUTO" {
		t.Fatalf("defaults not applied: %+v", cfg)
	}
}

func TestApplyDefaultsDuplicateModel(t *testing.T) {
	cfg := Config{Sources: []Source{
		{Name: "a", BaseURL: "http://x", Models: []Model{{ID: "m1"}}},
		{Name: "b", BaseURL: "http://y", Models: []Model{{ID: "m1"}}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatalf("duplicate model across sources should be allowed: %v", err)
	}
}

func TestStoreUpsertRemove(t *testing.T) {
	path := filepath.Join(t.TempDir(), "runtime.json")
	s := NewStore(path)
	if err := s.Load(); err != nil {
		t.Fatal(err)
	}
	if err := s.Upsert(Source{Name: "a", BaseURL: "http://a", APIKey: "sk-a", Models: []Model{{ID: "m"}}}); err != nil {
		t.Fatal(err)
	}
	if err := s.Upsert(Source{Name: "b", BaseURL: "http://b", Models: []Model{{ID: "m2"}}}); err != nil {
		t.Fatal(err)
	}
	if len(s.List()) != 2 {
		t.Fatalf("list = %d", len(s.List()))
	}
	removed, err := s.Remove("a")
	if err != nil || !removed {
		t.Fatalf("remove: %v %v", removed, err)
	}
	if len(s.List()) != 1 {
		t.Fatalf("after remove list = %d", len(s.List()))
	}
	// reload from disk
	s2 := NewStore(path)
	if err := s2.Load(); err != nil {
		t.Fatal(err)
	}
	if len(s2.List()) != 1 {
		t.Fatalf("reloaded list = %d", len(s2.List()))
	}
}

func TestStoreSecretEncryption(t *testing.T) {
	t.Setenv("LLMS_PROXY_MASTER_KEY", "")
	dir := t.TempDir()
	path := filepath.Join(dir, "runtime.json")
	s := NewStore(path)
	if s.box == nil {
		t.Fatal("expected secret box")
	}
	if err := s.Load(); err != nil {
		t.Fatal(err)
	}
	headers := map[string]string{"Authorization": "Bearer sk-hdr", "X-Custom": "plain"}
	if err := s.Upsert(Source{Name: "a", BaseURL: "http://a", APIKey: "sk-secret-123", Headers: headers}); err != nil {
		t.Fatal(err)
	}
	// file on disk must not contain plaintext secrets
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	for _, plain := range []string{"sk-secret-123", "Bearer sk-hdr"} {
		if strings.Contains(string(raw), plain) {
			t.Fatalf("secret %q stored in plaintext on disk", plain)
		}
	}
	// in-memory stays plaintext after the writes
	src := s.List()[0]
	if src.APIKey != "sk-secret-123" {
		t.Fatalf("in-memory api_key = %q", src.APIKey)
	}
	if src.Headers["Authorization"] != "Bearer sk-hdr" {
		t.Fatal("in-memory header not plaintext")
	}
	// reload: decrypted back
	s2 := NewStore(path)
	if err := s2.Load(); err != nil {
		t.Fatal(err)
	}
	if s2.List()[0].APIKey != "sk-secret-123" {
		t.Fatalf("reloaded api_key = %q", s2.List()[0].APIKey)
	}
}

func TestSecretBoxRoundTrip(t *testing.T) {
	dir := t.TempDir()
	box, err := NewSecretBox(filepath.Join(dir, "runtime.json"))
	if err != nil {
		t.Fatal(err)
	}
	v, err := box.Encrypt("sk-abc-xyz")
	if err != nil {
		t.Fatal(err)
	}
	if strings.HasPrefix(v, "sk-") || strings.Contains(v, "abc-xyz") {
		t.Fatalf("ciphertext leaked plaintext: %q", v)
	}
	out, err := box.Decrypt(v)
	if err != nil || out != "sk-abc-xyz" {
		t.Fatalf("roundtrip: %q %v", out, err)
	}
	if plain, err := box.Decrypt("sk-plain"); err != nil || plain != "sk-plain" {
		t.Fatalf("plain passthrough: %q %v", plain, err)
	}
	// wrong key must error
	bad, _ := NewSecretBox(filepath.Join(t.TempDir(), "runtime.json"))
	if _, err := bad.Decrypt(v); err == nil {
		t.Fatal("expected decrypt failure with wrong key")
	}
}
func TestRemoveSourceFromYAML(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cfg.yaml")
	content := `
listen: 127.0.0.1:9999
gateway_keys: [sk-1]
default_model: AUTO
adapter_dir: adapters
runtime_file: runtime.json
sources:
  - name: deepseek
    base_url: https://api.deepseek.com
    api_key: sk-d
    adapter: deepseek
    models:
      - id: deepseek-v4-flash
  - name: ollama
    base_url: http://127.0.0.1:11434
    adapter: ollama
    models:
      - id: llama3
`
	if err := os.WriteFile(path, []byte(content), 0644); err != nil {
		t.Fatal(err)
	}
	if err := RemoveSourceFromYAML(path, "deepseek"); err != nil {
		t.Fatalf("remove: %v", err)
	}
	reloaded, err := Load(path)
	if err != nil {
		t.Fatalf("reload: %v", err)
	}
	if len(reloaded.Sources) != 1 {
		t.Fatalf("sources = %d, want 1", len(reloaded.Sources))
	}
	if reloaded.Sources[0].Name != "ollama" {
		t.Fatalf("remaining = %q, want ollama", reloaded.Sources[0].Name)
	}
	if reloaded.Sources[0].Models[0].ID != "llama3" {
		t.Fatalf("remaining models broken: %+v", reloaded.Sources[0].Models)
	}
	// removing a non-existent name is a no-op that keeps the file valid
	if err := RemoveSourceFromYAML(path, "nope"); err != nil {
		t.Fatalf("remove missing: %v", err)
	}
	if cfg, err := Load(path); err != nil || len(cfg.Sources) != 1 {
		t.Fatalf("after no-op: %v %v", len(cfg.Sources), err)
	}
}

// TestSeedTemplatesOnce verifies preset templates are seeded exactly once:
// a user's edit survives a restart, and a deliberately deleted preset does
// not silently come back.
func TestSeedTemplatesOnce(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "runtime.json")
	presets := []SourceTemplate{
		{Name: "DeepSeek", BaseURL: "https://api.deepseek.com", Adapter: "deepseek",
			Models: []Model{{ID: "deepseek-chat", Kind: "chat"}}},
		{Name: "OpenAI", BaseURL: "https://api.openai.com/v1", Adapter: "openai",
			Models: []Model{{ID: "gpt-4o", Kind: "chat"}}},
	}

	// first run: both presets land in the store
	s1 := NewStore(path)
	if err := s1.Load(); err != nil {
		t.Fatal(err)
	}
	if err := s1.SeedTemplates(presets); err != nil {
		t.Fatalf("seed: %v", err)
	}
	if got := len(s1.ListTemplates()); got != 2 {
		t.Fatalf("after seed len = %d, want 2", got)
	}

	// user edits one preset, deletes the other
	edited := presets[0]
	edited.BaseURL = "https://my-proxy.internal/v1"
	if err := s1.UpsertTemplate(edited); err != nil {
		t.Fatal(err)
	}
	if _, err := s1.RemoveTemplate("OpenAI"); err != nil {
		t.Fatal(err)
	}

	// restart: re-seeding must be a no-op for both
	s2 := NewStore(path)
	if err := s2.Load(); err != nil {
		t.Fatal(err)
	}
	if err := s2.SeedTemplates(presets); err != nil {
		t.Fatalf("reseed: %v", err)
	}
	tpls := s2.ListTemplates()
	if len(tpls) != 1 {
		t.Fatalf("after restart len = %d, want 1 (deleted preset must not return): %+v", len(tpls), tpls)
	}
	if tpls[0].Name != "DeepSeek" || tpls[0].BaseURL != "https://my-proxy.internal/v1" {
		t.Fatalf("user edit lost on restart: %+v", tpls[0])
	}
}