ai-servers/llm-gateway/internal/provider/registry_test.go

package provider

import (
	"context"
	"io"
	"testing"

	"llm-gateway/internal/config"
)

// mockProvider implements the Provider interface for testing.
type mockProvider struct {
	name string
}

func (m *mockProvider) Name() string { return m.name }

func (m *mockProvider) ChatCompletion(_ context.Context, _ string, _ *ChatRequest) (*ChatResponse, error) {
	return nil, nil
}

func (m *mockProvider) ChatCompletionStream(_ context.Context, _ string, _ *ChatRequest) (io.ReadCloser, error) {
	return nil, nil
}

// newTestRegistry builds a Registry directly without going through config parsing.
func newTestRegistry(models []testModel) *Registry {
	r := &Registry{
		routes:    make(map[string][]Route),
		balancers: make(map[string]LoadBalancer),
		aliases:   make(map[string]string),
	}

	for _, m := range models {
		r.routes[m.name] = m.routes
		r.balancers[m.name] = &FirstBalancer{}
		r.order = append(r.order, m.name)
		for _, alias := range m.aliases {
			r.aliases[alias] = m.name
		}
	}

	return r
}

type testModel struct {
	name    string
	aliases []string
	routes  []Route
}

func TestRegistry_Lookup_Canonical(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name: "gpt-4",
			routes: []Route{
				{Provider: &mockProvider{name: "openai"}, ProviderModel: "gpt-4", Priority: 1},
			},
		},
	})

	routes, ok := reg.Lookup("gpt-4")
	if !ok {
		t.Fatal("expected Lookup to find gpt-4")
	}
	if len(routes) != 1 {
		t.Fatalf("expected 1 route, got %d", len(routes))
	}
	if routes[0].Provider.Name() != "openai" {
		t.Errorf("expected provider 'openai', got %q", routes[0].Provider.Name())
	}
}

func TestRegistry_Lookup_Alias(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name:    "gpt-4",
			aliases: []string{"gpt4", "gpt-4-latest"},
			routes: []Route{
				{Provider: &mockProvider{name: "openai"}, ProviderModel: "gpt-4", Priority: 1},
			},
		},
	})

	tests := []struct {
		name  string
		model string
		found bool
	}{
		{"canonical", "gpt-4", true},
		{"alias1", "gpt4", true},
		{"alias2", "gpt-4-latest", true},
		{"unknown", "gpt-5", false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			routes, ok := reg.Lookup(tt.model)
			if ok != tt.found {
				t.Fatalf("Lookup(%q) found=%v, want %v", tt.model, ok, tt.found)
			}
			if tt.found && len(routes) != 1 {
				t.Fatalf("expected 1 route, got %d", len(routes))
			}
		})
	}
}

func TestRegistry_ModelNames_IncludesAliases(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name:    "gpt-4",
			aliases: []string{"gpt4"},
			routes: []Route{
				{Provider: &mockProvider{name: "openai"}, ProviderModel: "gpt-4", Priority: 1},
			},
		},
		{
			name: "claude-3",
			routes: []Route{
				{Provider: &mockProvider{name: "anthropic"}, ProviderModel: "claude-3", Priority: 1},
			},
		},
	})

	names := reg.ModelNames()

	want := map[string]bool{"gpt-4": true, "gpt4": true, "claude-3": true}
	got := make(map[string]bool)
	for _, n := range names {
		got[n] = true
	}

	for name := range want {
		if !got[name] {
			t.Errorf("expected %q in ModelNames, not found", name)
		}
	}

	if len(names) != len(want) {
		t.Errorf("expected %d names, got %d: %v", len(want), len(names), names)
	}
}

func TestRegistry_AllRoutes_ShowsAliases(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name:    "gpt-4",
			aliases: []string{"gpt4", "gpt-4-latest"},
			routes: []Route{
				{Provider: &mockProvider{name: "openai"}, ProviderModel: "gpt-4", Priority: 1},
				{Provider: &mockProvider{name: "azure"}, ProviderModel: "gpt-4", Priority: 2},
			},
		},
	})

	allRoutes := reg.AllRoutes()
	if len(allRoutes) != 1 {
		t.Fatalf("expected 1 model, got %d", len(allRoutes))
	}

	m := allRoutes[0]
	if m.Name != "gpt-4" {
		t.Errorf("expected name 'gpt-4', got %q", m.Name)
	}

	aliasSet := make(map[string]bool)
	for _, a := range m.Aliases {
		aliasSet[a] = true
	}
	if !aliasSet["gpt4"] || !aliasSet["gpt-4-latest"] {
		t.Errorf("expected aliases [gpt4, gpt-4-latest], got %v", m.Aliases)
	}

	if len(m.Routes) != 2 {
		t.Fatalf("expected 2 routes, got %d", len(m.Routes))
	}
	if m.Routes[0].ProviderName != "openai" {
		t.Errorf("expected first route provider 'openai', got %q", m.Routes[0].ProviderName)
	}
	if m.Routes[1].ProviderName != "azure" {
		t.Errorf("expected second route provider 'azure', got %q", m.Routes[1].ProviderName)
	}
}

func TestRegistry_AllRoutes_ConfigOrder(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name: "model-b",
			routes: []Route{
				{Provider: &mockProvider{name: "prov"}, ProviderModel: "b", Priority: 1},
			},
		},
		{
			name: "model-a",
			routes: []Route{
				{Provider: &mockProvider{name: "prov"}, ProviderModel: "a", Priority: 1},
			},
		},
	})

	allRoutes := reg.AllRoutes()
	if len(allRoutes) != 2 {
		t.Fatalf("expected 2 models, got %d", len(allRoutes))
	}
	if allRoutes[0].Name != "model-b" {
		t.Errorf("expected first model 'model-b', got %q", allRoutes[0].Name)
	}
	if allRoutes[1].Name != "model-a" {
		t.Errorf("expected second model 'model-a', got %q", allRoutes[1].Name)
	}
}

func TestRegistry_PrioritySorting(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name: "multi-provider",
			routes: []Route{
				{Provider: &mockProvider{name: "low-priority"}, ProviderModel: "m", Priority: 3},
				{Provider: &mockProvider{name: "high-priority"}, ProviderModel: "m", Priority: 1},
				{Provider: &mockProvider{name: "mid-priority"}, ProviderModel: "m", Priority: 2},
			},
		},
	})

	// Note: routes are stored as given (sorting happens during buildFromConfig).
	// For this test we verify AllRoutes returns them in stored order.
	allRoutes := reg.AllRoutes()
	if len(allRoutes) != 1 {
		t.Fatalf("expected 1 model, got %d", len(allRoutes))
	}

	routes := allRoutes[0].Routes
	if len(routes) != 3 {
		t.Fatalf("expected 3 routes, got %d", len(routes))
	}

	// Verify the priorities are present
	priorities := make(map[int]bool)
	for _, r := range routes {
		priorities[r.Priority] = true
	}
	for _, p := range []int{1, 2, 3} {
		if !priorities[p] {
			t.Errorf("expected priority %d in routes", p)
		}
	}
}

func TestRegistry_NewRegistry_UnknownProvider(t *testing.T) {
	cfg := &config.Config{
		Models: []config.ModelConfig{
			{
				Name: "test-model",
				Routes: []config.RouteConfig{
					{Provider: "nonexistent", Model: "m"},
				},
			},
		},
	}

	_, err := NewRegistry(cfg)
	if err == nil {
		t.Fatal("expected error for unknown provider, got nil")
	}
}

func TestRegistry_Lookup_NotFound(t *testing.T) {
	reg := newTestRegistry([]testModel{
		{
			name: "gpt-4",
			routes: []Route{
				{Provider: &mockProvider{name: "openai"}, ProviderModel: "gpt-4", Priority: 1},
			},
		},
	})

	_, ok := reg.Lookup("nonexistent")
	if ok {
		t.Fatal("expected Lookup to return false for nonexistent model")
	}
}