package router_test import ( "context" "errors" "net/http" "net/http/httptest" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/zap" "github.com/veylant/ia-gateway/internal/config" "github.com/veylant/ia-gateway/internal/middleware" "github.com/veylant/ia-gateway/internal/provider" "github.com/veylant/ia-gateway/internal/router" "github.com/veylant/ia-gateway/internal/routing" ) // ─── Mock adapter ──────────────────────────────────────────────────────────── // mockAdapter records which methods were called and returns preset responses. type mockAdapter struct { name string sendCalled bool streamCalled bool healthErr error } func (m *mockAdapter) Send(_ context.Context, req *provider.ChatRequest) (*provider.ChatResponse, error) { m.sendCalled = true return &provider.ChatResponse{ ID: "mock-" + m.name, Model: req.Model, Choices: []provider.Choice{{ Message: provider.Message{Role: "assistant", Content: "response from " + m.name}, }}, }, nil } func (m *mockAdapter) Stream(_ context.Context, _ *provider.ChatRequest, w http.ResponseWriter) error { m.streamCalled = true _, _ = w.Write([]byte("data: [DONE]\n\n")) return nil } func (m *mockAdapter) Validate(req *provider.ChatRequest) error { if req.Model == "" { return errors.New("model required") } return nil } func (m *mockAdapter) HealthCheck(_ context.Context) error { return m.healthErr } // ─── Helpers ───────────────────────────────────────────────────────────────── func newAdapters() (map[string]provider.Adapter, map[string]*mockAdapter) { mocks := map[string]*mockAdapter{ "openai": {name: "openai"}, "anthropic": {name: "anthropic"}, "mistral": {name: "mistral"}, "ollama": {name: "ollama"}, } adapters := make(map[string]provider.Adapter, len(mocks)) for k, v := range mocks { adapters[k] = v } return adapters, mocks } func adminCtx() context.Context { return middleware.WithClaims(context.Background(), &middleware.UserClaims{ UserID: "u1", TenantID: "t1", Roles: []string{"admin"}, }) } func userCtx(allowedModels ...string) context.Context { return middleware.WithClaims(context.Background(), &middleware.UserClaims{ UserID: "u2", TenantID: "t1", Roles: []string{"user"}, }) } func auditorCtx() context.Context { return middleware.WithClaims(context.Background(), &middleware.UserClaims{ UserID: "u3", TenantID: "t1", Roles: []string{"auditor"}, }) } func rbacConfig() *config.RBACConfig { return &config.RBACConfig{ UserAllowedModels: []string{"gpt-4o-mini", "gpt-3.5-turbo", "mistral-small"}, AuditorCanComplete: false, } } // ─── Routing ───────────────────────────────────────────────────────────────── func TestRouter_Send_RoutesToOpenAI(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["openai"].sendCalled, "openai adapter should have been called") assert.False(t, mocks["anthropic"].sendCalled) } func TestRouter_Send_RoutesToAnthropic(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "claude-3-5-sonnet-20241022", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["anthropic"].sendCalled, "anthropic adapter should have been called") assert.False(t, mocks["openai"].sendCalled) } func TestRouter_Send_RoutesToMistral(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "mistral-medium", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["mistral"].sendCalled) } func TestRouter_Send_RoutesToOllama_Llama(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "llama3", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["ollama"].sendCalled) } func TestRouter_Send_UnknownModel_FallsBackToOpenAI(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "some-unknown-model", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["openai"].sendCalled, "unknown model should fall back to openai") } // ─── RBAC ──────────────────────────────────────────────────────────────────── func TestRouter_Send_AdminCanAccessAll(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) for _, model := range []string{"gpt-4o", "claude-3-opus", "mistral-medium", "llama3"} { _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: model, Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) assert.NoError(t, err, "admin should access %q", model) } } func TestRouter_Send_UserAllowedModel_Routed(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(userCtx(), &provider.ChatRequest{ Model: "gpt-4o-mini", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["openai"].sendCalled) } func TestRouter_Send_UserBlockedUnauthorizedModel(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(userCtx(), &provider.ChatRequest{ Model: "claude-3-opus", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.Error(t, err) assert.Contains(t, err.Error(), "claude-3-opus") } func TestRouter_Send_AuditorBlocked(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) _, err := r.Send(auditorCtx(), &provider.ChatRequest{ Model: "gpt-4o-mini", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.Error(t, err) assert.Contains(t, err.Error(), "auditor") } // ─── Stream ────────────────────────────────────────────────────────────────── func TestRouter_Stream_RoutesToAnthropic(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) rec := httptest.NewRecorder() err := r.Stream(adminCtx(), &provider.ChatRequest{ Model: "claude-3-5-sonnet-20241022", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }, rec) require.NoError(t, err) assert.True(t, mocks["anthropic"].streamCalled) } // ─── HealthCheck ───────────────────────────────────────────────────────────── func TestRouter_HealthCheck_AllHealthy(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) err := r.HealthCheck(context.Background()) require.NoError(t, err) } func TestRouter_HealthCheck_OneUnhealthy(t *testing.T) { adapters, mocks := newAdapters() mocks["anthropic"].healthErr = errors.New("connection refused") r := router.New(adapters, rbacConfig(), zap.NewNop()) err := r.HealthCheck(context.Background()) require.Error(t, err) assert.Contains(t, err.Error(), "anthropic") } // ─── Validate ──────────────────────────────────────────────────────────────── func TestRouter_Validate_DelegatestoAdapter(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) // openai adapter requires model; should delegate and return its error err := r.Validate(&provider.ChatRequest{Model: ""}) require.Error(t, err) } func TestRouter_Validate_ValidRequest(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) err := r.Validate(&provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) } // ─── Stream RBAC ───────────────────────────────────────────────────────────── func TestRouter_Stream_AuditorBlocked(t *testing.T) { adapters, _ := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) rec := httptest.NewRecorder() err := r.Stream(auditorCtx(), &provider.ChatRequest{ Model: "gpt-4o-mini", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }, rec) require.Error(t, err) assert.Contains(t, err.Error(), "auditor") } // ─── No fallback configured ─────────────────────────────────────────────────── func TestRouter_Send_NoAdapters_ReturnsError(t *testing.T) { r := router.New(map[string]provider.Adapter{}, rbacConfig(), zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.Error(t, err) } // ─── Engine integration ─────────────────────────────────────────────────────── // failAdapter always returns an error from Send. type failAdapter struct{ *mockAdapter } func (f *failAdapter) Send(_ context.Context, _ *provider.ChatRequest) (*provider.ChatResponse, error) { return nil, errors.New("provider unavailable") } func newEngineWithRule(rule routing.RoutingRule) *routing.Engine { store := routing.NewMemStore() _, _ = store.Create(context.Background(), rule) return routing.New(store, 30*time.Second, zap.NewNop()) } // TestRouter_WithEngine_SensitivityRoutesToOllama verifies that when the engine // matches a critical-sensitivity rule, the request goes to ollama not openai. func TestRouter_WithEngine_SensitivityRoutesToOllama(t *testing.T) { adapters, mocks := newAdapters() engine := newEngineWithRule(routing.RoutingRule{ TenantID: "t1", Name: "critical → ollama", Priority: 10, IsEnabled: true, Conditions: []routing.Condition{ {Field: "request.sensitivity", Operator: "gte", Value: "high"}, }, Action: routing.Action{Provider: "ollama"}, }) r := router.NewWithEngine(adapters, rbacConfig(), engine, zap.NewNop()) ctx := routing.WithSensitivity(adminCtx(), routing.SensitivityCritical) _, err := r.Send(ctx, &provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["ollama"].sendCalled, "critical sensitivity should route to ollama") assert.False(t, mocks["openai"].sendCalled) } // TestRouter_WithEngine_NoMatch_FallsBackToStaticRules verifies that when no // engine rule matches, static prefix routing is used (backward compat). func TestRouter_WithEngine_NoMatch_FallsBackToStaticRules(t *testing.T) { adapters, mocks := newAdapters() // Engine rule requires "finance" dept — request has no dept. engine := newEngineWithRule(routing.RoutingRule{ TenantID: "t1", Name: "finance only", Priority: 10, IsEnabled: true, Conditions: []routing.Condition{ {Field: "user.department", Operator: "eq", Value: "finance"}, }, Action: routing.Action{Provider: "ollama"}, }) r := router.NewWithEngine(adapters, rbacConfig(), engine, zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["openai"].sendCalled, "no engine match → static prefix rule → openai") assert.False(t, mocks["ollama"].sendCalled) } // TestRouter_FallbackChain_PrimaryFails_TriesNext verifies that when the primary // provider fails, the next provider in the fallback chain is tried. func TestRouter_FallbackChain_PrimaryFails_TriesNext(t *testing.T) { mocks := map[string]*mockAdapter{ "openai": {name: "openai"}, "anthropic": {name: "anthropic"}, "mistral": {name: "mistral"}, "ollama": {name: "ollama"}, } adapters := make(map[string]provider.Adapter, len(mocks)) for k, v := range mocks { adapters[k] = v } // Replace ollama with a failing adapter. adapters["ollama"] = &failAdapter{mocks["ollama"]} engine := newEngineWithRule(routing.RoutingRule{ TenantID: "t1", Name: "ollama with openai fallback", Priority: 10, IsEnabled: true, Conditions: []routing.Condition{}, Action: routing.Action{Provider: "ollama", FallbackProviders: []string{"openai"}}, }) r := router.NewWithEngine(adapters, rbacConfig(), engine, zap.NewNop()) _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "gpt-4o", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["openai"].sendCalled, "fallback to openai after ollama fails") } // TestRouter_NilEngine_BackwardCompat verifies that a router without an engine // behaves exactly as before Sprint 5 (static prefix rules). func TestRouter_NilEngine_BackwardCompat(t *testing.T) { adapters, mocks := newAdapters() r := router.New(adapters, rbacConfig(), zap.NewNop()) // no engine _, err := r.Send(adminCtx(), &provider.ChatRequest{ Model: "claude-3-opus", Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["anthropic"].sendCalled, "claude- prefix → anthropic") } // TestRouter_WithEngine_DepartmentRouting verifies department-based routing. func TestRouter_WithEngine_DepartmentRouting(t *testing.T) { adapters, mocks := newAdapters() engine := newEngineWithRule(routing.RoutingRule{ TenantID: "t1", Name: "engineering → anthropic", Priority: 10, IsEnabled: true, Conditions: []routing.Condition{ {Field: "user.department", Operator: "eq", Value: "engineering"}, }, Action: routing.Action{Provider: "anthropic"}, }) r := router.NewWithEngine(adapters, rbacConfig(), engine, zap.NewNop()) ctx := middleware.WithClaims(context.Background(), &middleware.UserClaims{ UserID: "u1", TenantID: "t1", Roles: []string{"admin"}, Department: "engineering", }) _, err := r.Send(ctx, &provider.ChatRequest{ Model: "gpt-4o", // model prefix → openai, but engine should override Messages: []provider.Message{{Role: "user", Content: "hi"}}, }) require.NoError(t, err) assert.True(t, mocks["anthropic"].sendCalled, "engineering dept should route to anthropic") assert.False(t, mocks["openai"].sendCalled) }