veylant/internal/routing/engine_test.go

package routing_test

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"

	"github.com/veylant/ia-gateway/internal/routing"
)

const tenantA = "tenant-a"
const tenantB = "tenant-b"

// newEngine creates an Engine with a MemStore pre-seeded with the given rules.
func newEngine(rules ...routing.RoutingRule) *routing.Engine {
	store := routing.NewMemStore()
	for _, r := range rules {
		_, _ = store.Create(context.Background(), r)
	}
	return routing.New(store, 30*time.Second, zap.NewNop())
}

func rule(tenantID, name string, priority int, action routing.Action, conds ...routing.Condition) routing.RoutingRule {
	return routing.RoutingRule{
		TenantID:   tenantID,
		Name:       name,
		Priority:   priority,
		IsEnabled:  true,
		Action:     action,
		Conditions: conds,
	}
}

func actionFor(provider string, fallbacks ...string) routing.Action {
	return routing.Action{Provider: provider, FallbackProviders: fallbacks}
}

// ─── Basic match ─────────────────────────────────────────────────────────────

func TestEngine_Evaluate_NoRules_ReturnsFalse(t *testing.T) {
	e := newEngine()
	_, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	require.NoError(t, err)
	assert.False(t, matched)
}

func TestEngine_Evaluate_CatchAll_Matches(t *testing.T) {
	e := newEngine(rule(tenantA, "catch-all", 9999, actionFor("openai")))
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "openai", action.Provider)
}

func TestEngine_Evaluate_NoMatchForTenant_ReturnsFalse(t *testing.T) {
	e := newEngine(rule(tenantA, "r1", 10, actionFor("openai")))
	_, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantB})
	require.NoError(t, err)
	assert.False(t, matched)
}

// ─── Sensitivity routing ──────────────────────────────────────────────────────

func TestEngine_CriticalSensitivity_RoutesToOllama(t *testing.T) {
	r := rule(tenantA, "critical → ollama", 10, actionFor("ollama", "openai"),
		routing.Condition{Field: "request.sensitivity", Operator: "gte", Value: "high"},
	)
	e := newEngine(r)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Sensitivity: routing.SensitivityCritical,
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "ollama", action.Provider)
	assert.Equal(t, []string{"openai"}, action.FallbackProviders)
}

func TestEngine_LowSensitivity_DoesNotMatchHighRule(t *testing.T) {
	r := rule(tenantA, "high → ollama", 10, actionFor("ollama"),
		routing.Condition{Field: "request.sensitivity", Operator: "gte", Value: "high"},
	)
	e := newEngine(r)
	_, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Sensitivity: routing.SensitivityLow,
	})
	require.NoError(t, err)
	assert.False(t, matched)
}

func TestEngine_MediumSensitivity_MatchesMediumRule(t *testing.T) {
	r := rule(tenantA, "medium+", 10, actionFor("mistral"),
		routing.Condition{Field: "request.sensitivity", Operator: "gte", Value: "medium"},
	)
	e := newEngine(r)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Sensitivity: routing.SensitivityMedium,
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "mistral", action.Provider)
}

// ─── Department routing ───────────────────────────────────────────────────────

func TestEngine_Department_Finance_RoutesToOllama(t *testing.T) {
	r := rule(tenantA, "finance → ollama", 10, actionFor("ollama"),
		routing.Condition{Field: "user.department", Operator: "eq", Value: "finance"},
	)
	e := newEngine(r)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:   tenantA,
		Department: "finance",
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "ollama", action.Provider)
}

func TestEngine_Department_Engineering_RoutesToAnthropic(t *testing.T) {
	r := rule(tenantA, "eng → anthropic", 10, actionFor("anthropic"),
		routing.Condition{Field: "user.department", Operator: "eq", Value: "engineering"},
	)
	e := newEngine(r)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:   tenantA,
		Department: "engineering",
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "anthropic", action.Provider)
}

func TestEngine_Department_RH_OtherDept_DoesNotMatch(t *testing.T) {
	r := rule(tenantA, "rh → ollama", 10, actionFor("ollama"),
		routing.Condition{Field: "user.department", Operator: "eq", Value: "rh"},
	)
	e := newEngine(r)
	_, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:   tenantA,
		Department: "marketing",
	})
	require.NoError(t, err)
	assert.False(t, matched)
}

// ─── Role routing ─────────────────────────────────────────────────────────────

func TestEngine_Role_Admin_MatchesAdminRule(t *testing.T) {
	r := rule(tenantA, "admin → anthropic", 10, actionFor("anthropic"),
		routing.Condition{Field: "user.role", Operator: "in", Value: []interface{}{"admin", "manager"}},
	)
	e := newEngine(r)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID: tenantA,
		UserRole: "admin",
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "anthropic", action.Provider)
}

// ─── Priority ordering ────────────────────────────────────────────────────────

func TestEngine_Priority_LowestWins(t *testing.T) {
	// Two catch-all rules — priority 10 should win over priority 20
	r1 := rule(tenantA, "catch-all-low", 20, actionFor("openai"))
	r2 := rule(tenantA, "catch-all-high", 10, actionFor("anthropic"))
	e := newEngine(r1, r2)
	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "anthropic", action.Provider, "priority 10 should win over priority 20")
}

func TestEngine_Priority_SpecificBeforeCatchAll(t *testing.T) {
	specific := rule(tenantA, "critical → ollama", 5, actionFor("ollama"),
		routing.Condition{Field: "request.sensitivity", Operator: "gte", Value: "high"},
	)
	catchAll := rule(tenantA, "catch-all", 9999, actionFor("openai"))
	e := newEngine(catchAll, specific) // insert out of order; engine should sort

	action, matched, err := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Sensitivity: routing.SensitivityCritical,
	})
	require.NoError(t, err)
	assert.True(t, matched)
	assert.Equal(t, "ollama", action.Provider, "specific rule should take priority over catch-all")
}

// ─── Multiple conditions (AND logic) ─────────────────────────────────────────

func TestEngine_MultipleConditions_AllMustMatch(t *testing.T) {
	r := rule(tenantA, "finance + high sens", 10, actionFor("ollama"),
		routing.Condition{Field: "user.department", Operator: "eq", Value: "finance"},
		routing.Condition{Field: "request.sensitivity", Operator: "gte", Value: "high"},
	)
	e := newEngine(r)

	// Only department matches, not sensitivity
	_, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Department:  "finance",
		Sensitivity: routing.SensitivityLow,
	})
	assert.False(t, matched, "both conditions must match")

	// Both match
	action, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:    tenantA,
		Department:  "finance",
		Sensitivity: routing.SensitivityHigh,
	})
	assert.True(t, matched)
	assert.Equal(t, "ollama", action.Provider)
}

// ─── Token estimate ───────────────────────────────────────────────────────────

func TestEngine_TokenEstimate_LargePrompt_RoutesToOllama(t *testing.T) {
	r := rule(tenantA, "big prompts → ollama", 10, actionFor("ollama"),
		routing.Condition{Field: "request.token_estimate", Operator: "gte", Value: float64(4000)},
	)
	e := newEngine(r)

	_, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:      tenantA,
		TokenEstimate: 100,
	})
	assert.False(t, matched)

	action, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID:      tenantA,
		TokenEstimate: 5000,
	})
	assert.True(t, matched)
	assert.Equal(t, "ollama", action.Provider)
}

// ─── Model routing ────────────────────────────────────────────────────────────

func TestEngine_Model_MatchesPrefix(t *testing.T) {
	r := rule(tenantA, "gpt-4 → azure", 10, actionFor("azure"),
		routing.Condition{Field: "request.model", Operator: "matches", Value: "gpt-4"},
	)
	e := newEngine(r)
	action, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{
		TenantID: tenantA,
		Model:    "gpt-4o-mini",
	})
	assert.True(t, matched)
	assert.Equal(t, "azure", action.Provider)
}

// ─── Disabled rules ───────────────────────────────────────────────────────────

func TestEngine_DisabledRule_NotEvaluated(t *testing.T) {
	store := routing.NewMemStore()
	r, _ := store.Create(context.Background(), routing.RoutingRule{
		TenantID:   tenantA,
		Name:       "disabled",
		Priority:   10,
		IsEnabled:  false, // disabled
		Conditions: []routing.Condition{},
		Action:     routing.Action{Provider: "anthropic"},
	})
	_ = r

	e := routing.New(store, 30*time.Second, zap.NewNop())
	_, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	assert.False(t, matched, "disabled rule must not be evaluated")
}

// ─── Fallback providers ───────────────────────────────────────────────────────

func TestEngine_Action_FallbackProvidersPreserved(t *testing.T) {
	r := rule(tenantA, "with fallbacks", 10,
		routing.Action{Provider: "anthropic", FallbackProviders: []string{"openai", "mistral"}},
	)
	e := newEngine(r)
	action, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	assert.True(t, matched)
	assert.Equal(t, "anthropic", action.Provider)
	assert.Equal(t, []string{"openai", "mistral"}, action.FallbackProviders)
}

// ─── Model override ───────────────────────────────────────────────────────────

func TestEngine_Action_ModelOverride(t *testing.T) {
	r := rule(tenantA, "model override", 10,
		routing.Action{Provider: "ollama", Model: "llama3"},
	)
	e := newEngine(r)
	action, matched, _ := e.Evaluate(context.Background(), &routing.RoutingContext{TenantID: tenantA})
	assert.True(t, matched)
	assert.Equal(t, "llama3", action.Model)
}

// ─── Benchmark ───────────────────────────────────────────────────────────────

func BenchmarkEngine_Evaluate_10Rules(b *testing.B) {
	rules := make([]routing.RoutingRule, 10)
	for i := range rules {
		rules[i] = routing.RoutingRule{
			TenantID:  tenantA,
			Name:      "rule",
			Priority:  (i + 1) * 10,
			IsEnabled: true,
			Action:    routing.Action{Provider: "openai"},
			Conditions: []routing.Condition{
				{Field: "user.department", Operator: "eq", Value: "nonexistent"},
			},
		}
	}
	// Last rule is catch-all to ensure we always match
	rules[9].Conditions = []routing.Condition{}

	e := newEngine(rules...)
	rctx := &routing.RoutingContext{TenantID: tenantA, Department: "finance"}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _, _ = e.Evaluate(context.Background(), rctx)
	}
}