bifrost/plugins/governance/store_test.go

package governance

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	bifrost "github.com/maximhq/bifrost/core"
	"github.com/maximhq/bifrost/core/schemas"
	"github.com/maximhq/bifrost/framework/configstore"
	configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestGovernanceStore_GetVirtualKey tests lock-free VK retrieval
func TestGovernanceStore_GetVirtualKey(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{
			*buildVirtualKey("vk1", "sk-bf-test1", "Test VK 1", true),
			*buildVirtualKey("vk2", "sk-bf-test2", "Test VK 2", false),
		},
	}, nil)
	require.NoError(t, err)

	tests := []struct {
		name    string
		vkValue string
		wantNil bool
		wantID  string
	}{
		{
			name:    "Found active VK",
			vkValue: "sk-bf-test1",
			wantNil: false,
			wantID:  "vk1",
		},
		{
			name:    "Found inactive VK",
			vkValue: "sk-bf-test2",
			wantNil: false,
			wantID:  "vk2",
		},
		{
			name:    "VK not found",
			vkValue: "sk-bf-nonexistent",
			wantNil: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			vk, exists := store.GetVirtualKey(context.Background(), tt.vkValue)
			if tt.wantNil {
				assert.False(t, exists)
				assert.Nil(t, vk)
			} else {
				assert.True(t, exists)
				assert.NotNil(t, vk)
				assert.Equal(t, tt.wantID, vk.ID)
			}
		})
	}
}

// TestGovernanceStore_ConcurrentReads tests lock-free concurrent reads
func TestGovernanceStore_ConcurrentReads(t *testing.T) {
	logger := NewMockLogger()
	vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", true)
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
	}, nil)
	require.NoError(t, err)

	// Launch 100 concurrent readers
	var wg sync.WaitGroup
	readCount := atomic.Int64{}
	errorCount := atomic.Int64{}

	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 100; j++ {
				vk, exists := store.GetVirtualKey(context.Background(), "sk-bf-test")
				if !exists || vk == nil {
					errorCount.Add(1)
					return
				}
				readCount.Add(1)
			}
		}()
	}

	wg.Wait()

	assert.Equal(t, int64(10000), readCount.Load(), "Expected 10000 successful reads")
	assert.Equal(t, int64(0), errorCount.Load(), "Expected 0 errors")
}

// TestGovernanceStore_CheckBudget_SingleBudget tests budget validation with single budget
func TestGovernanceStore_CheckBudget_SingleBudget(t *testing.T) {
	logger := NewMockLogger()
	budget := buildBudgetWithUsage("budget1", 100.0, 50.0, "1d")
	vk := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", budget)

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*budget},
	}, nil)
	require.NoError(t, err)

	// Retrieve VK with budget
	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")

	tests := []struct {
		name      string
		usage     float64
		maxLimit  float64
		shouldErr bool
	}{
		{
			name:      "Usage below limit",
			usage:     50.0,
			maxLimit:  100.0,
			shouldErr: false,
		},
		{
			name:      "Usage at limit (should fail)",
			usage:     100.0,
			maxLimit:  100.0,
			shouldErr: true,
		},
		{
			name:      "Usage exceeds limit",
			usage:     150.0,
			maxLimit:  100.0,
			shouldErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create new budget with test usage
			testBudget := buildBudgetWithUsage("budget1", tt.maxLimit, tt.usage, "1d")
			testVK := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", testBudget)
			testStore, _ := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
				VirtualKeys: []configstoreTables.TableVirtualKey{*testVK},
				Budgets:     []configstoreTables.TableBudget{*testBudget},
			}, nil)

			testVK, _ = testStore.GetVirtualKey(context.Background(), "sk-bf-test")
			_, err := testStore.CheckVirtualKeyBudget(context.Background(), testVK, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
			if tt.shouldErr {
				assert.Error(t, err, "Expected error for usage check")
			} else {
				assert.NoError(t, err, "Expected no error for usage check")
			}
		})
	}
}

// TestGovernanceStore_CheckBudget_HierarchyValidation tests multi-level budget hierarchy
func TestGovernanceStore_CheckBudget_HierarchyValidation(t *testing.T) {
	logger := NewMockLogger()

	// Create budgets at different levels
	vkBudget := buildBudgetWithUsage("vk-budget", 100.0, 50.0, "1d")
	teamBudget := buildBudgetWithUsage("team-budget", 500.0, 200.0, "1d")
	customerBudget := buildBudgetWithUsage("customer-budget", 1000.0, 400.0, "1d")

	// Build hierarchy
	team := buildTeam("team1", "Team 1", teamBudget)
	customer := buildCustomer("customer1", "Customer 1", customerBudget)
	team.CustomerID = &customer.ID
	team.Customer = customer

	vk := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", vkBudget)
	vk.TeamID = &team.ID
	vk.Team = team

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*vkBudget, *teamBudget, *customerBudget},
		Teams:       []configstoreTables.TableTeam{*team},
		Customers:   []configstoreTables.TableCustomer{*customer},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")

	// Test: All budgets under limit should pass
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	assert.NoError(t, err, "Should pass when all budgets are under limit")

	// Test: If VK budget exceeds limit, should fail
	// Update the budget directly in the budgets map (since UpdateVirtualKeyInMemory preserves usage)
	if len(vk.Budgets) > 0 {
		budgetID := vk.Budgets[0].ID
		if budgetValue, exists := store.budgets.Load(budgetID); exists && budgetValue != nil {
			if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
				budget.CurrentUsage = 100.0
				store.budgets.Store(budgetID, budget)
			}
		}
	}
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail when VK budget exceeds limit")
}

// TestGovernanceStore_MultiBudget_AllUnderLimit tests that requests pass when all budgets are under their limits
func TestGovernanceStore_MultiBudget_AllUnderLimit(t *testing.T) {
	logger := NewMockLogger()

	// Create VK with hourly ($10) and daily ($100) budgets
	hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 5.0, "1h")
	dailyBudget := buildBudgetWithUsage("daily", 100.0, 40.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	// Add provider config so the resolver allows the provider
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	assert.NoError(t, err, "Should pass when all budgets are under limit")
}

// TestGovernanceStore_MultiBudget_SmallBudgetExceeded tests that request is blocked when the smaller budget exceeds its limit
func TestGovernanceStore_MultiBudget_SmallBudgetExceeded(t *testing.T) {
	logger := NewMockLogger()

	// Hourly at limit, daily still has room
	hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 10.0, "1h")
	dailyBudget := buildBudgetWithUsage("daily", 100.0, 40.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine")
	assert.Contains(t, err.Error(), "budget exceeded")
}

// TestGovernanceStore_MultiBudget_LargeBudgetExceeded tests that request is blocked when only the larger budget exceeds
func TestGovernanceStore_MultiBudget_LargeBudgetExceeded(t *testing.T) {
	logger := NewMockLogger()

	// Hourly has room, but daily is at limit
	hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 3.0, "1h")
	dailyBudget := buildBudgetWithUsage("daily", 100.0, 100.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail when daily budget is exceeded even though hourly is fine")
	assert.Contains(t, err.Error(), "budget exceeded")
}

// TestGovernanceStore_MultiBudget_UsageUpdatesAllBudgets tests that usage updates are applied to every budget in the hierarchy
func TestGovernanceStore_MultiBudget_UsageUpdatesAllBudgets(t *testing.T) {
	logger := NewMockLogger()

	hourlyBudget := buildBudget("hourly", 10.0, "1h")
	dailyBudget := buildBudget("daily", 100.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")

	// Simulate a $3.50 request
	err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 3.50)
	require.NoError(t, err)

	// Both budgets should reflect the cost
	hourlyVal, exists := store.budgets.Load("hourly")
	require.True(t, exists)
	assert.InDelta(t, 3.50, hourlyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Hourly budget should reflect usage")

	dailyVal, exists := store.budgets.Load("daily")
	require.True(t, exists)
	assert.InDelta(t, 3.50, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Daily budget should reflect usage")

	// Second request: $7.00 — should push hourly over limit
	err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 7.00)
	require.NoError(t, err)

	hourlyVal, _ = store.budgets.Load("hourly")
	assert.InDelta(t, 10.50, hourlyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Hourly budget should accumulate")

	dailyVal, _ = store.budgets.Load("daily")
	assert.InDelta(t, 10.50, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Daily budget should accumulate")

	// Now CheckBudget should fail (hourly exceeded)
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail after usage exceeds hourly budget")
	assert.Contains(t, err.Error(), "budget exceeded")
}

// TestGovernanceStore_MultiBudget_ProviderConfigBudgets tests that provider-config-level multi-budgets are enforced
func TestGovernanceStore_MultiBudget_ProviderConfigBudgets(t *testing.T) {
	logger := NewMockLogger()

	// Provider-level budgets: hourly $5 (exceeded), daily $50 (ok)
	pcHourly := buildBudgetWithUsage("pc-hourly", 5.0, 5.0, "1h")
	pcDaily := buildBudgetWithUsage("pc-daily", 50.0, 10.0, "1d")

	pc := buildProviderConfigWithBudgets("openai", []string{"*"},
		[]configstoreTables.TableBudget{*pcHourly, *pcDaily})

	vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableVirtualKeyProviderConfig{pc})

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*pcHourly, *pcDaily},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail when provider config hourly budget is exceeded")
	assert.Contains(t, err.Error(), "budget exceeded")
}

// TestGovernanceStore_MultiBudget_VKAndProviderConfigCombined tests budgets at both VK and provider config levels
func TestGovernanceStore_MultiBudget_VKAndProviderConfigCombined(t *testing.T) {
	logger := NewMockLogger()

	// VK-level budgets: all under limit
	vkMonthly := buildBudgetWithUsage("vk-monthly", 1000.0, 200.0, "1M")

	// Provider-config-level budgets: hourly at limit
	pcHourly := buildBudgetWithUsage("pc-hourly", 5.0, 5.0, "1h")

	pc := buildProviderConfigWithBudgets("openai", []string{"*"},
		[]configstoreTables.TableBudget{*pcHourly})

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*vkMonthly})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{pc}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*vkMonthly, *pcHourly},
	}, nil)
	require.NoError(t, err)

	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")

	// Provider config budget exceeded → should block even though VK budget is fine
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	require.Error(t, err, "Should fail: provider config budget exceeded even though VK budget is fine")
	assert.Contains(t, err.Error(), "budget exceeded")
}

// TestGovernanceStore_MultiBudget_ResolverBlocksOnBudgetExceeded tests that the full resolver flow blocks when any budget is exceeded
func TestGovernanceStore_MultiBudget_ResolverBlocksOnBudgetExceeded(t *testing.T) {
	logger := NewMockLogger()

	// Two VK-level budgets: hourly at limit, daily has room
	hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 10.0, "1h")
	dailyBudget := buildBudgetWithUsage("daily", 100.0, 30.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	resolver := NewBudgetResolver(store, nil, logger, nil)
	ctx := &schemas.BifrostContext{}

	result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
	assertDecision(t, DecisionBudgetExceeded, result)
	assert.Contains(t, result.Reason, "budget exceeded")
}

// TestGovernanceStore_MultiBudget_ResolverAllowsUnderLimit tests that the full resolver flow allows requests when all budgets are under limit
func TestGovernanceStore_MultiBudget_ResolverAllowsUnderLimit(t *testing.T) {
	logger := NewMockLogger()

	hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 5.0, "1h")
	dailyBudget := buildBudgetWithUsage("daily", 100.0, 30.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	resolver := NewBudgetResolver(store, nil, logger, nil)
	ctx := &schemas.BifrostContext{}

	result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
	assertDecision(t, DecisionAllow, result)
}

// TestGovernanceStore_MultiBudget_UsageDrivesBlockAfterRequests tests the full lifecycle:
// start under limit → accumulate usage → eventually hit a budget → get blocked
func TestGovernanceStore_MultiBudget_UsageDrivesBlockAfterRequests(t *testing.T) {
	logger := NewMockLogger()

	// Tight hourly ($2), generous daily ($100)
	hourlyBudget := buildBudget("hourly", 2.0, "1h")
	dailyBudget := buildBudget("daily", 100.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
	}, nil)
	require.NoError(t, err)

	resolver := NewBudgetResolver(store, nil, logger, nil)

	// Request 1: $0.80 — both budgets fine
	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
	require.NoError(t, err)

	ctx := &schemas.BifrostContext{}
	result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
	assertDecision(t, DecisionAllow, result)

	// Request 2: $0.80 — still fine ($1.60 total)
	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
	require.NoError(t, err)

	ctx = &schemas.BifrostContext{}
	result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
	assertDecision(t, DecisionAllow, result)

	// Request 3: $0.80 — pushes hourly to $2.40 > $2.00 limit → blocked
	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
	require.NoError(t, err)

	ctx = &schemas.BifrostContext{}
	result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
	assertDecision(t, DecisionBudgetExceeded, result)
	assert.Contains(t, result.Reason, "budget exceeded")

	// Verify daily budget is still under limit
	dailyVal, exists := store.budgets.Load("daily")
	require.True(t, exists)
	assert.InDelta(t, 2.40, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01,
		"Daily budget should be at $2.40, well under $100 limit")
}

// TestGovernanceStore_MultiBudget_CalendarAligned tests that calendar-aligned budgets are stored and retrievable
func TestGovernanceStore_MultiBudget_CalendarAligned(t *testing.T) {
	logger := NewMockLogger()

	// Calendar alignment is a VK-level setting — budgets don't have it
	dailyBudget := &configstoreTables.TableBudget{
		ID:            "daily-cal",
		MaxLimit:      50.0,
		CurrentUsage:  10.0,
		ResetDuration: "1d",
		LastReset:     time.Now(),
	}
	monthlyBudget := &configstoreTables.TableBudget{
		ID:            "monthly-cal",
		MaxLimit:      1000.0,
		CurrentUsage:  200.0,
		ResetDuration: "1M",
		LastReset:     time.Now(),
	}

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*dailyBudget, *monthlyBudget})
	vk.CalendarAligned = true // VK-level setting applies to all budgets
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*dailyBudget, *monthlyBudget},
	}, nil)
	require.NoError(t, err)

	// Verify VK-level calendar_aligned is set
	vk, _ = store.GetVirtualKey(context.Background(), "sk-bf-test")
	assert.True(t, vk.CalendarAligned, "VK should have calendar_aligned=true")

	// Both under limit — should pass
	_, err = store.CheckVirtualKeyBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
	assert.NoError(t, err)
}

// TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete tests CreateVirtualKeyInMemory and DeleteVirtualKeyInMemory
// properly store and clean up multi-budget entries
func TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete(t *testing.T) {
	logger := NewMockLogger()

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	b1 := buildBudget("b1", 10.0, "1h")
	b2 := buildBudget("b2", 100.0, "1d")

	vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
		[]configstoreTables.TableBudget{*b1, *b2})
	vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
		buildProviderConfig("openai", []string{"*"}),
	}

	// Create
	store.CreateVirtualKeyInMemory(context.Background(), vk)

	_, exists := store.budgets.Load("b1")
	assert.True(t, exists, "Budget b1 should be in memory after create")
	_, exists = store.budgets.Load("b2")
	assert.True(t, exists, "Budget b2 should be in memory after create")

	retrieved, found := store.GetVirtualKey(context.Background(), "sk-bf-test")
	require.True(t, found)
	assert.Len(t, retrieved.Budgets, 2, "VK should have 2 budgets")

	// Delete
	store.DeleteVirtualKeyInMemory(context.Background(), "vk1")

	_, exists = store.budgets.Load("b1")
	assert.False(t, exists, "Budget b1 should be removed after delete")
	_, exists = store.budgets.Load("b2")
	assert.False(t, exists, "Budget b2 should be removed after delete")

	_, found = store.GetVirtualKey(context.Background(), "sk-bf-test")
	assert.False(t, found, "VK should not be found after delete")
}

// TestGovernanceStore_UpdateRateLimitUsage_TokensAndRequests tests atomic rate limit usage updates
func TestGovernanceStore_UpdateRateLimitUsage_TokensAndRequests(t *testing.T) {
	logger := NewMockLogger()

	rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 0)
	vk := buildVirtualKeyWithRateLimit("vk1", "sk-bf-test", "Test VK", rateLimit)

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		RateLimits:  []configstoreTables.TableRateLimit{*rateLimit},
	}, nil)
	require.NoError(t, err)

	// Test updating tokens
	err = store.UpdateVirtualKeyRateLimitUsageInMemory(context.Background(), vk, schemas.OpenAI, 500, true, false)
	assert.NoError(t, err, "Rate limit update should succeed")

	// Retrieve the updated rate limit from the main RateLimits map
	governanceData := store.GetGovernanceData(context.Background())
	updatedRateLimit, exists := governanceData.RateLimits["rl1"]
	require.True(t, exists, "Rate limit should exist")
	require.NotNil(t, updatedRateLimit)

	assert.Equal(t, int64(500), updatedRateLimit.TokenCurrentUsage, "Token usage should be updated")
	assert.Equal(t, int64(0), updatedRateLimit.RequestCurrentUsage, "Request usage should not change")

	// Test updating requests
	err = store.UpdateVirtualKeyRateLimitUsageInMemory(context.Background(), vk, schemas.OpenAI, 0, false, true)
	assert.NoError(t, err, "Rate limit update should succeed")

	// Retrieve the updated rate limit again
	governanceData = store.GetGovernanceData(context.Background())
	updatedRateLimit, exists = governanceData.RateLimits["rl1"]
	require.True(t, exists, "Rate limit should exist")
	require.NotNil(t, updatedRateLimit)

	assert.Equal(t, int64(500), updatedRateLimit.TokenCurrentUsage, "Token usage should not change")
	assert.Equal(t, int64(1), updatedRateLimit.RequestCurrentUsage, "Request usage should be incremented")
}

// TestGovernanceStore_ResetExpiredRateLimits tests rate limit reset
func TestGovernanceStore_ResetExpiredRateLimits(t *testing.T) {
	logger := NewMockLogger()

	// Create rate limit that's already expired
	duration := "1m"
	rateLimit := &configstoreTables.TableRateLimit{
		ID:                   "rl1",
		TokenMaxLimit:        ptrInt64(10000),
		TokenCurrentUsage:    5000,
		TokenResetDuration:   &duration,
		TokenLastReset:       time.Now().Add(-2 * time.Minute), // Expired
		RequestMaxLimit:      ptrInt64(1000),
		RequestCurrentUsage:  500,
		RequestResetDuration: &duration,
		RequestLastReset:     time.Now().Add(-2 * time.Minute), // Expired
	}

	vk := buildVirtualKeyWithRateLimit("vk1", "sk-bf-test", "Test VK", rateLimit)

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		RateLimits:  []configstoreTables.TableRateLimit{*rateLimit},
	}, nil)
	require.NoError(t, err)

	// Reset expired rate limits
	expiredRateLimits := store.ResetExpiredRateLimitsInMemory(context.Background())
	err = store.ResetExpiredRateLimits(context.Background(), expiredRateLimits)
	assert.NoError(t, err, "Reset should succeed")

	// Retrieve the updated VK to check rate limit changes
	updatedVK, _ := store.GetVirtualKey(context.Background(), "sk-bf-test")
	require.NotNil(t, updatedVK)
	require.NotNil(t, updatedVK.RateLimit)

	assert.Equal(t, int64(0), updatedVK.RateLimit.TokenCurrentUsage, "Token usage should be reset")
	assert.Equal(t, int64(0), updatedVK.RateLimit.RequestCurrentUsage, "Request usage should be reset")
}

// TestGovernanceStore_ResetExpiredBudgets tests budget reset
func TestGovernanceStore_ResetExpiredBudgets(t *testing.T) {
	logger := NewMockLogger()

	// Create budget that's already expired
	budget := &configstoreTables.TableBudget{
		ID:            "budget1",
		MaxLimit:      100.0,
		CurrentUsage:  75.0,
		ResetDuration: "1d",
		LastReset:     time.Now().Add(-48 * time.Hour), // Expired
	}

	vk := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", budget)

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
		Budgets:     []configstoreTables.TableBudget{*budget},
	}, nil)
	require.NoError(t, err)

	// Reset expired budgets
	expiredBudgets := store.ResetExpiredBudgetsInMemory(context.Background())
	err = store.ResetExpiredBudgets(context.Background(), expiredBudgets)
	assert.NoError(t, err, "Reset should succeed")

	// Retrieve the updated VK to check budget changes
	updatedVK, _ := store.GetVirtualKey(context.Background(), "sk-bf-test")
	require.NotNil(t, updatedVK)
	require.True(t, len(updatedVK.Budgets) > 0, "VK should have budgets")

	assert.Equal(t, 0.0, updatedVK.Budgets[0].CurrentUsage, "Budget usage should be reset")
}

// TestGovernanceStore_GetAllBudgets tests retrieving all budgets
func TestGovernanceStore_GetAllBudgets(t *testing.T) {
	logger := NewMockLogger()

	budgets := []configstoreTables.TableBudget{
		*buildBudget("budget1", 100.0, "1d"),
		*buildBudget("budget2", 500.0, "1d"),
		*buildBudget("budget3", 1000.0, "1d"),
	}

	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
		Budgets: budgets,
	}, nil)
	require.NoError(t, err)

	allBudgets := store.GetGovernanceData(context.Background()).Budgets
	assert.Equal(t, 3, len(allBudgets), "Should have 3 budgets")
	assert.NotNil(t, allBudgets["budget1"])
	assert.NotNil(t, allBudgets["budget2"])
	assert.NotNil(t, allBudgets["budget3"])
}

// TestGovernanceStore_RoutingRules_CreateAndRetrieve tests creating and retrieving routing rules
func TestGovernanceStore_RoutingRules_CreateAndRetrieve(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	// Create a global routing rule
	rule1 := &configstoreTables.TableRoutingRule{
		ID:            "1",
		Name:          "Global Rule",
		Description:   "Test global routing rule",
		Enabled:       true,
		CelExpression: "model == 'gpt-4o'",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai"), Model: bifrost.Ptr("gpt-4"), Weight: 1.0},
		},
		Fallbacks:       nil,
		ParsedFallbacks: []string{"azure/gpt-4-turbo"},
		Scope:           "global",
		ScopeID:         nil,
		Priority:        10,
		CreatedAt:       time.Now(),
		UpdatedAt:       time.Now(),
	}

	// Create a team-scoped routing rule
	teamID := "team-123"
	rule2 := &configstoreTables.TableRoutingRule{
		ID:            "2",
		Name:          "Team Rule",
		Description:   "Test team routing rule",
		Enabled:       true,
		CelExpression: "model in ['gpt-4o', 'gpt-4-turbo']",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("azure"), Weight: 1.0},
		},
		Fallbacks:       nil,
		ParsedFallbacks: []string{"groq/mixtral-8x7b"},
		Scope:           "team",
		ScopeID:         &teamID,
		Priority:        20,
		CreatedAt:       time.Now(),
		UpdatedAt:       time.Now(),
	}

	// Store rules in memory
	err = store.UpdateRoutingRuleInMemory(context.Background(), rule1)
	require.NoError(t, err)
	err = store.UpdateRoutingRuleInMemory(context.Background(), rule2)
	require.NoError(t, err)

	// Test retrieval by scope
	globalRules := store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 1, len(globalRules))
	assert.Equal(t, "Global Rule", globalRules[0].Name)

	teamRules := store.GetScopedRoutingRules(context.Background(), "team", teamID)
	assert.Equal(t, 1, len(teamRules))
	assert.Equal(t, "Team Rule", teamRules[0].Name)

	// Test ListRoutingRules
	allRules := store.GetAllRoutingRules(context.Background())
	assert.Equal(t, 2, len(allRules))
}

// TestGovernanceStore_RoutingRules_PriorityOrdering tests that rules are sorted by priority
func TestGovernanceStore_RoutingRules_PriorityOrdering(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	// Create rules with different priorities
	rules := []*configstoreTables.TableRoutingRule{
		{
			ID:       "1",
			Name:     "Priority 5",
			Priority: 5,
			Scope:    "global",
			ScopeID:  nil,
			Enabled:  true,
		},
		{
			ID:       "2",
			Name:     "Priority 20",
			Priority: 20,
			Scope:    "global",
			ScopeID:  nil,
			Enabled:  true,
		},
		{
			ID:       "3",
			Name:     "Priority 10",
			Priority: 10,
			Scope:    "global",
			ScopeID:  nil,
			Enabled:  true,
		},
	}

	for _, rule := range rules {
		err := store.UpdateRoutingRuleInMemory(context.Background(), rule)
		require.NoError(t, err)
	}

	// Retrieve and verify ordering (sorted by priority ASC, so lower numbers first)
	retrieved := store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 3, len(retrieved))
	assert.Equal(t, 5, retrieved[0].Priority)
	assert.Equal(t, 10, retrieved[1].Priority)
	assert.Equal(t, 20, retrieved[2].Priority)
}

// TestGovernanceStore_RoutingRules_DisabledRulesFiltered tests that disabled rules are filtered out
func TestGovernanceStore_RoutingRules_DisabledRulesFiltered(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	enabledRule := &configstoreTables.TableRoutingRule{
		ID:      "1",
		Name:    "Enabled Rule",
		Enabled: true,
		Scope:   "global",
		ScopeID: nil,
	}

	disabledRule := &configstoreTables.TableRoutingRule{
		ID:      "2",
		Name:    "Disabled Rule",
		Enabled: false,
		Scope:   "global",
		ScopeID: nil,
	}

	err = store.UpdateRoutingRuleInMemory(context.Background(), enabledRule)
	require.NoError(t, err)
	err = store.UpdateRoutingRuleInMemory(context.Background(), disabledRule)
	require.NoError(t, err)

	// Only enabled rules should be returned
	retrieved := store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 1, len(retrieved))
	assert.Equal(t, "Enabled Rule", retrieved[0].Name)
}

// TestGovernanceStore_RoutingRules_DeleteRule tests deleting a routing rule
func TestGovernanceStore_RoutingRules_DeleteRule(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:      "1",
		Name:    "Test Rule",
		Enabled: true,
		Scope:   "global",
		ScopeID: nil,
	}

	// Add rule
	err = store.UpdateRoutingRuleInMemory(context.Background(), rule)
	require.NoError(t, err)

	retrieved := store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 1, len(retrieved))

	// Delete rule
	err = store.DeleteRoutingRuleInMemory(context.Background(), rule.ID)
	require.NoError(t, err)

	// Verify deletion
	retrieved = store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 0, len(retrieved))
}

// TestGovernanceStore_RateLimitStatus tests rate limit status calculation
func TestGovernanceStore_RateLimitStatus(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	// Create a rate limit with 1000 token limit
	limit := int64(1000)
	rateLimitID := "provider:openai:ratelimit"
	rl := &configstoreTables.TableRateLimit{
		ID:                rateLimitID,
		TokenMaxLimit:     &limit,
		TokenCurrentUsage: 500,
	}

	store.rateLimits.Store(rateLimitID, rl)

	// Create a provider config that references the rate limit
	providerConfig := &configstoreTables.TableProvider{
		Name:        "openai",
		RateLimitID: &rateLimitID,
	}
	store.providers.Store("openai", providerConfig)

	// Get status
	status := store.GetBudgetAndRateLimitStatus(context.Background(), "", schemas.ModelProvider("openai"), nil, nil, nil, nil)

	assert.NotNil(t, status)
	assert.Equal(t, 50.0, status.RateLimitTokenPercentUsed)

	// Update usage to exhausted state
	rl.TokenCurrentUsage = 1000
	status = store.GetBudgetAndRateLimitStatus(context.Background(), "", schemas.ModelProvider("openai"), nil, nil, nil, nil)

	assert.Equal(t, 100.0, status.RateLimitTokenPercentUsed)
}

// TestGovernanceStore_BudgetStatus tests budget status calculation
func TestGovernanceStore_BudgetStatus(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	budgetID := "provider:openai:budget"
	budget := &configstoreTables.TableBudget{
		ID:           budgetID,
		MaxLimit:     100.0,
		CurrentUsage: 60.0,
	}

	store.budgets.Store(budgetID, budget)

	// Create a provider config that references the budget
	providerConfig := &configstoreTables.TableProvider{
		Name:     "openai",
		BudgetID: &budgetID,
	}
	store.providers.Store("openai", providerConfig)

	// Get status
	status := store.GetBudgetAndRateLimitStatus(context.Background(), "", schemas.ModelProvider("openai"), nil, nil, nil, nil)

	assert.NotNil(t, status)
	assert.Equal(t, 60.0, status.BudgetPercentUsed)

	// Update usage to exhausted state
	budget.CurrentUsage = 100.0
	status = store.GetBudgetAndRateLimitStatus(context.Background(), "", schemas.ModelProvider("openai"), nil, nil, nil, nil)

	assert.Equal(t, 100.0, status.BudgetPercentUsed)
}

// TestGovernanceStore_RoutingRules_MultipleScopes tests rules with multiple scopes
func TestGovernanceStore_RoutingRules_MultipleScopes(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	customerID := "cust-123"
	teamID := "team-456"

	// Create rules for different scopes
	globalRule := &configstoreTables.TableRoutingRule{
		ID: "1", Name: "Global", Scope: "global", ScopeID: nil, Priority: 10, Enabled: true,
	}
	customerRule := &configstoreTables.TableRoutingRule{
		ID: "2", Name: "Customer", Scope: "customer", ScopeID: &customerID, Priority: 20, Enabled: true,
	}
	teamRule := &configstoreTables.TableRoutingRule{
		ID: "3", Name: "Team", Scope: "team", ScopeID: &teamID, Priority: 30, Enabled: true,
	}

	require.NoError(t, store.UpdateRoutingRuleInMemory(context.Background(), globalRule))
	require.NoError(t, store.UpdateRoutingRuleInMemory(context.Background(), customerRule))
	require.NoError(t, store.UpdateRoutingRuleInMemory(context.Background(), teamRule))

	// Test global scope
	globalRules := store.GetScopedRoutingRules(context.Background(), "global", "")
	assert.Equal(t, 1, len(globalRules))
	assert.Equal(t, "Global", globalRules[0].Name)

	// Test customer scope
	custRules := store.GetScopedRoutingRules(context.Background(), "customer", customerID)
	assert.Equal(t, 1, len(custRules))
	assert.Equal(t, "Customer", custRules[0].Name)

	// Test team scope
	teamRules := store.GetScopedRoutingRules(context.Background(), "team", teamID)
	assert.Equal(t, 1, len(teamRules))
	assert.Equal(t, "Team", teamRules[0].Name)

	// ListAll should return all rules sorted by priority ASC (lower numbers = higher priority)
	allRules := store.GetAllRoutingRules(context.Background())
	assert.Equal(t, 3, len(allRules))
	assert.Equal(t, 10, allRules[0].Priority) // Global (highest)
	assert.Equal(t, 20, allRules[1].Priority) // Customer
	assert.Equal(t, 30, allRules[2].Priority) // Team (lowest)
}

// TestCompileAndCacheProgram tests CEL program compilation and caching
func TestCompileAndCacheProgram(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:            "rule-1",
		Name:          "Test Rule",
		CelExpression: "model == 'gpt-4o' && tokens_used < 80.0",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai")},
		},
		Enabled: true,
	}

	// First compilation
	program1, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program1)

	// Verify it's cached - second call should return cached program
	program2, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program2)

	// Both should be the same cached instance
	assert.Equal(t, program1, program2)
}

// TestCompileAndCacheProgram_InvalidExpression tests error handling for invalid CEL
func TestCompileAndCacheProgram_InvalidExpression(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:            "rule-invalid",
		Name:          "Invalid Rule",
		CelExpression: "model == gpt-4o'", // Syntax error
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai")},
		},
		Enabled: true,
	}

	_, err = store.GetRoutingProgram(context.Background(), rule)
	assert.Error(t, err)

	// Invalid rule should not be cached - attempting to get it again should fail
	_, err = store.GetRoutingProgram(context.Background(), rule)
	assert.Error(t, err)
}

// TestCompileAndCacheProgram_CacheInvalidation tests cache invalidation on rule update
func TestCompileAndCacheProgram_CacheInvalidation(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:            "rule-update",
		Name:          "Update Rule",
		CelExpression: "model == 'gpt-4o'",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai")},
		},
		Enabled: true,
		Scope:   "global",
	}

	// Compile and cache
	program1, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program1)

	// Update rule in memory (should invalidate cache)
	rule.CelExpression = "model == 'gpt-4-turbo'"
	err = store.UpdateRoutingRuleInMemory(context.Background(), rule)
	require.NoError(t, err)

	// Recompile should work
	program2, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program2)
}

// TestCompileAndCacheProgram_CacheInvalidationOnDelete tests cache invalidation on rule deletion
func TestCompileAndCacheProgram_CacheInvalidationOnDelete(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:            "rule-delete",
		Name:          "Delete Rule",
		CelExpression: "provider == 'openai'",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai")},
		},
		Enabled: true,
		Scope:   "global",
	}

	// Compile and cache
	_, err = store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)

	// Delete rule (should invalidate cache)
	err = store.DeleteRoutingRuleInMemory(context.Background(), rule.ID)
	require.NoError(t, err)

	// After deletion, we can't verify cache directly, but the rule is gone from storage
}

// TestCompileAndCacheProgram_EmptyExpression tests compilation of empty CEL expression (defaults to "true")
func TestCompileAndCacheProgram_EmptyExpression(t *testing.T) {
	logger := NewMockLogger()
	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
	require.NoError(t, err)

	rule := &configstoreTables.TableRoutingRule{
		ID:            "rule-empty",
		Name:          "Empty Rule",
		CelExpression: "",
		Targets: []configstoreTables.TableRoutingTarget{
			{Provider: bifrost.Ptr("openai")},
		},
		Enabled: true,
	}

	program, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program)

	// Verify caching works - second call should return same program
	program2, err := store.GetRoutingProgram(context.Background(), rule)
	require.NoError(t, err)
	assert.NotNil(t, program2)
	assert.Equal(t, program, program2)
}

// Utility functions for tests
func ptrInt64(i int64) *int64 {
	return &i
}