package governance import ( "context" "testing" "time" "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" ) // ============================================================================ // Store Tests - Provider Budget // ============================================================================ func TestStore_CheckProviderBudget_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.NoError(t, err, "Should allow when no provider config exists") } func TestStore_CheckProviderBudget_NoBudget(t *testing.T) { logger := NewMockLogger() provider := buildProviderWithGovernance("openai", nil, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, }, nil) require.NoError(t, err) _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.NoError(t, err, "Should allow when provider has no budget") } func TestStore_CheckProviderBudget_WithinLimit(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.NoError(t, err, "Should allow when budget is within limit") } func TestStore_CheckProviderBudget_Exceeded(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.Error(t, err, "Should reject when budget is exceeded") assert.Contains(t, err.Error(), "budget exceeded") } func TestStore_CheckProviderBudget_WithBaseline(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 90.0, "1h") provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) // With baseline that would exceed limit baselines := map[string]float64{"budget1": 15.0} _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, baselines) assert.Error(t, err, "Should reject when current usage + baseline exceeds limit") assert.Contains(t, err.Error(), "budget exceeded") } // ============================================================================ // Store Tests - Provider Rate Limit // ============================================================================ func TestStore_CheckProviderRateLimit_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.NoError(t, err, "Should allow when no provider config exists") assert.Equal(t, DecisionAllow, decision) } func TestStore_CheckProviderRateLimit_NoRateLimit(t *testing.T) { logger := NewMockLogger() provider := buildProviderWithGovernance("openai", nil, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, }, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.NoError(t, err, "Should allow when provider has no rate limit") assert.Equal(t, DecisionAllow, decision) } func TestStore_CheckProviderRateLimit_TokenLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.Error(t, err, "Should reject when provider token limit is exceeded") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_CheckProviderRateLimit_RequestLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.Error(t, err, "Should reject when provider request limit is exceeded") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } func TestStore_CheckProviderRateLimit_BothLimitsExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 1000) // Both at max provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.Error(t, err, "Should reject when both provider token and request limits are exceeded") assert.Equal(t, DecisionRateLimited, decision) // General rate limited when both are exceeded assert.Contains(t, err.Error(), "rate limit") } func TestStore_CheckProviderRateLimit_WithinLimits(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 5000, 1000, 500) // Both within limits provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.NoError(t, err, "Should allow when provider rate limits are within limits") assert.Equal(t, DecisionAllow, decision) } // ============================================================================ // Store Tests - Model Budget // ============================================================================ func TestStore_CheckModelBudget_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.NoError(t, err, "Should allow when no model config exists") } func TestStore_CheckModelBudget_ModelOnly_WithinLimit(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.NoError(t, err, "Should allow when model budget is within limit") } func TestStore_CheckModelBudget_ModelOnly_Exceeded(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.Error(t, err, "Should reject when model budget is exceeded") assert.Contains(t, err.Error(), "budget exceeded") } func TestStore_CheckModelBudget_ModelWithProvider_WithinLimit(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4", &providerStr, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.NoError(t, err, "Should allow when model+provider budget is within limit") } func TestStore_CheckModelBudget_ModelWithProvider_Exceeded(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4", &providerStr, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.Error(t, err, "Should reject when model+provider budget is exceeded") assert.Contains(t, err.Error(), "budget exceeded") } func TestStore_CheckModelBudget_BothModelAndModelProvider_ChecksBoth(t *testing.T) { logger := NewMockLogger() // Model-only budget (exceeded) budget1 := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, budget1, nil) // Model+provider budget (within limit) budget2 := buildBudget("budget2", 200.0, "1h") providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, budget2, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, Budgets: []configstoreTables.TableBudget{*budget1, *budget2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.Error(t, err, "Should reject when model-only budget is exceeded, even if model+provider budget is OK") assert.Contains(t, err.Error(), "budget exceeded") } func TestStore_CheckModelBudget_ProviderSpecific_DifferentProvider_Passes(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has budget (exceeded) budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) // Request with Azure (different provider) for same model should pass provider := schemas.Azure _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: provider}, nil) assert.NoError(t, err, "Should allow when model config is provider-specific and different provider is used") } // ============================================================================ // Store Tests - Model Rate Limit // ============================================================================ func TestStore_CheckModelRateLimit_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when no model config exists") assert.Equal(t, DecisionAllow, decision) } func TestStore_CheckModelRateLimit_ModelOnly_TokenLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model token limit is exceeded") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_CheckModelRateLimit_ModelOnly_RequestLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model request limit is exceeded") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } func TestStore_CheckModelRateLimit_ModelWithProvider_WithinLimits(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 5000, 1000, 500) // Within limits providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when model+provider rate limits are within limits") assert.Equal(t, DecisionAllow, decision) } func TestStore_CheckModelRateLimit_BothModelAndModelProvider_ChecksBoth(t *testing.T) { logger := NewMockLogger() // Model-only rate limit (exceeded) rateLimit1 := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit1) // Model+provider rate limit (within limit) rateLimit2 := buildRateLimitWithUsage("rl2", 20000, 5000, 2000, 500) providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, nil, rateLimit2) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit1, *rateLimit2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model-only rate limit is exceeded") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_CheckModelRateLimit_BothModelAndModelProvider_ChecksBoth_RequestLimit(t *testing.T) { logger := NewMockLogger() // Model-only rate limit (request limit exceeded) rateLimit1 := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit1) // Model+provider rate limit (within limit) rateLimit2 := buildRateLimitWithUsage("rl2", 20000, 5000, 2000, 500) providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, nil, rateLimit2) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit1, *rateLimit2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model-only rate limit (request limit) is exceeded") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } func TestStore_CheckModelRateLimit_ProviderSpecific_DifferentProvider_Passes(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) // Request with Azure (different provider) for same model should pass provider := schemas.Azure decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when model config is provider-specific and different provider is used") assert.Equal(t, DecisionAllow, decision) } func TestStore_CheckModelRateLimit_ProviderSpecific_DifferentProvider_Passes_RequestLimit(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (request limit exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) // Request with Azure (different provider) for same model should pass provider := schemas.Azure decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when model config is provider-specific and different provider is used (request limit)") assert.Equal(t, DecisionAllow, decision) } // ============================================================================ // Store Tests - Update Provider Budget Usage // ============================================================================ func TestStore_UpdateProviderBudgetUsage_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "", schemas.OpenAI, 10.0) assert.NoError(t, err, "Should not error when no provider config exists") } func TestStore_UpdateProviderBudgetUsage_UpdatesUsage(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "", schemas.OpenAI, 10.0) assert.NoError(t, err, "Should successfully update provider budget usage") // Verify usage was updated _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.NoError(t, err, "Should still be within limit after first update") // Update again to exceed err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "", schemas.OpenAI, 95.0) assert.NoError(t, err, "Should successfully update provider budget usage even when exceeding") // Now should be exceeded _, err = store.CheckProviderBudget(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil) assert.Error(t, err, "Should be exceeded after second update") assert.Contains(t, err.Error(), "budget exceeded") } // ============================================================================ // Store Tests - Update Provider Rate Limit Usage // ============================================================================ func TestStore_UpdateProviderRateLimitUsage_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "", schemas.OpenAI, 1000, true, true) assert.NoError(t, err, "Should not error when no provider config exists") } func TestStore_UpdateProviderRateLimitUsage_UpdatesTokens(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimit("rl1", 10000, 1000) provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "", schemas.OpenAI, 5000, true, false) assert.NoError(t, err, "Should successfully update provider token usage") // Check that tokens were updated but requests were not decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.NoError(t, err, "Should still be within token limit") assert.Equal(t, DecisionAllow, decision) // Update tokens to exceed err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "", schemas.OpenAI, 6000, true, false) assert.NoError(t, err, "Should successfully update provider token usage even when exceeding") // Now should be exceeded decision, err = store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.Error(t, err, "Should reject when provider token limit is exceeded after update") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_UpdateProviderRateLimitUsage_UpdatesRequests(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimit("rl1", 10000, 1000) provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) // Update requests 500 times for i := 0; i < 500; i++ { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "", schemas.OpenAI, 0, false, true) assert.NoError(t, err, "Should successfully update provider request usage") } // Should still be within limit decision, err := store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.NoError(t, err, "Should allow when provider request limit is within limit") assert.Equal(t, DecisionAllow, decision) // Update 500 more times to exceed for i := 0; i < 500; i++ { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "", schemas.OpenAI, 0, false, true) assert.NoError(t, err, "Should successfully update provider request usage even when exceeding") } // Now should be exceeded decision, err = store.CheckProviderRateLimit(context.Background(), &EvaluationRequest{Provider: schemas.OpenAI}, nil, nil) assert.Error(t, err, "Should reject when provider request limit is exceeded after update") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } // ============================================================================ // Store Tests - Update Model Budget Usage // ============================================================================ func TestStore_UpdateModelBudgetUsage_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "gpt-4", provider, 10.0) assert.NoError(t, err, "Should not error when no model config exists") } func TestStore_UpdateModelBudgetUsage_ModelOnly_UpdatesUsage(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "gpt-4", provider, 10.0) assert.NoError(t, err, "Should successfully update model budget usage") // Verify usage was updated _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.NoError(t, err, "Should still be within limit after first update") // Update again to exceed err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "gpt-4", provider, 95.0) assert.NoError(t, err, "Should successfully update model budget usage even when exceeding") // Now should be exceeded _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.Error(t, err, "Should be exceeded after second update") assert.Contains(t, err.Error(), "budget exceeded") } func TestStore_UpdateModelBudgetUsage_ModelWithProvider_UpdatesBoth(t *testing.T) { logger := NewMockLogger() // Model-only budget budget1 := buildBudget("budget1", 100.0, "1h") modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, budget1, nil) // Model+provider budget budget2 := buildBudget("budget2", 200.0, "1h") providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, budget2, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, Budgets: []configstoreTables.TableBudget{*budget1, *budget2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "gpt-4", provider, 10.0) assert.NoError(t, err, "Should successfully update both model-only and model+provider budget usage") // Both budgets should be updated // Check model-only budget _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.NoError(t, err, "Should still be within limit") // Update to exceed model-only budget err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "gpt-4", provider, 95.0) assert.NoError(t, err, "Should successfully update model budget usage even when exceeding") // Now model-only budget should be exceeded _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil) assert.Error(t, err, "Should be exceeded when model-only budget is exceeded") assert.Contains(t, err.Error(), "budget exceeded") } // ============================================================================ // Store Tests - Update Model Rate Limit Usage // ============================================================================ func TestStore_UpdateModelRateLimitUsage_NoConfig(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 1000, true, true) assert.NoError(t, err, "Should not error when no model config exists") } func TestStore_UpdateModelRateLimitUsage_ModelOnly_UpdatesUsage(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimit("rl1", 10000, 1000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 5000, true, false) assert.NoError(t, err, "Should successfully update model token usage") // Should still be within limit decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when model token limit is within limit") assert.Equal(t, DecisionAllow, decision) // Update to exceed err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 6000, true, false) assert.NoError(t, err, "Should successfully update model token usage even when exceeding") // Now should be exceeded decision, err = store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model token limit is exceeded after update") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_UpdateModelRateLimitUsage_ModelWithProvider_UpdatesUsage(t *testing.T) { logger := NewMockLogger() // Model-only rate limit rateLimit1 := buildRateLimit("rl1", 10000, 1000) modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit1) // Model+provider rate limit rateLimit2 := buildRateLimit("rl2", 20000, 2000) providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, nil, rateLimit2) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit1, *rateLimit2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 5000, true, false) assert.NoError(t, err, "Should successfully update both model-only and model+provider token usage") // Should still be within limit decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when both rate limits are within limit") assert.Equal(t, DecisionAllow, decision) // Update to exceed model-only rate limit (should fail at model-only level) err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 6000, true, false) assert.NoError(t, err, "Should successfully update model token usage even when exceeding") // Now should be exceeded (model-only rate limit exceeded) decision, err = store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model-only token limit is exceeded after update") assert.Equal(t, DecisionTokenLimited, decision) assert.Contains(t, err.Error(), "token limit exceeded") } func TestStore_UpdateModelRateLimitUsage_ModelOnly_UpdatesUsage_RequestLimit(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimit("rl1", 10000, 1000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) provider := schemas.OpenAI // Update requests 500 times for range 500 { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 0, false, true) assert.NoError(t, err, "Should successfully update model request usage") } // Should still be within limit decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when model request limit is within limit") assert.Equal(t, DecisionAllow, decision) // Update 500 more times to exceed for range 500 { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 0, false, true) assert.NoError(t, err, "Should successfully update model request usage even when exceeding") } // Now should be exceeded decision, err = store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model request limit is exceeded after update") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } func TestStore_UpdateModelRateLimitUsage_ModelWithProvider_UpdatesUsage_RequestLimit(t *testing.T) { logger := NewMockLogger() // Model-only rate limit rateLimit1 := buildRateLimit("rl1", 10000, 1000) modelConfig1 := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit1) // Model+provider rate limit rateLimit2 := buildRateLimit("rl2", 20000, 2000) providerStr := "openai" modelConfig2 := buildModelConfig("mc2", "gpt-4", &providerStr, nil, rateLimit2) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig1, *modelConfig2}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit1, *rateLimit2}, }, nil) require.NoError(t, err) provider := schemas.OpenAI // Update requests 500 times (should update both model-only and model+provider) for range 500 { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 0, false, true) assert.NoError(t, err, "Should successfully update both model-only and model+provider request usage") } // Should still be within limit decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.NoError(t, err, "Should allow when both rate limits are within limit") assert.Equal(t, DecisionAllow, decision) // Update 500 more times to exceed model-only rate limit for range 500 { err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4", provider, 0, false, true) assert.NoError(t, err, "Should successfully update model request usage even when exceeding") } // Now should be exceeded (model-only rate limit exceeded) decision, err = store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: provider}, nil, nil) assert.Error(t, err, "Should reject when model-only request limit is exceeded after update") assert.Equal(t, DecisionRequestLimited, decision) assert.Contains(t, err.Error(), "request limit exceeded") } // ============================================================================ // Resolver Tests - EvaluateModelAndProviderRequest // ============================================================================ func TestResolver_EvaluateModelAndProviderRequest_NoConfigs(t *testing.T) { logger := NewMockLogger() store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionAllow, result) } func TestResolver_EvaluateModelAndProviderRequest_ProviderBudgetExceeded(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionBudgetExceeded, result) assert.Contains(t, result.Reason, "Provider-level budget exceeded") } func TestResolver_EvaluateModelAndProviderRequest_ProviderRateLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionTokenLimited, result) assert.Contains(t, result.Reason, "Provider-level rate limit check failed") } func TestResolver_EvaluateModelAndProviderRequest_ModelBudgetExceeded(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionBudgetExceeded, result) assert.Contains(t, result.Reason, "Model-level budget exceeded") } func TestResolver_EvaluateModelAndProviderRequest_ModelRateLimitExceeded(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionTokenLimited, result) assert.Contains(t, result.Reason, "Model-level rate limit check failed") } func TestResolver_EvaluateModelAndProviderRequest_ModelRateLimitExceeded_RequestLimit(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionRequestLimited, result) assert.Contains(t, result.Reason, "Model-level rate limit check failed") } func TestResolver_EvaluateModelAndProviderRequest_ProviderBudgetThenModelBudget(t *testing.T) { logger := NewMockLogger() // Provider budget exceeded providerBudget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") provider := buildProviderWithGovernance("openai", providerBudget, nil) // Model budget within limit modelBudget := buildBudget("budget2", 200.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, modelBudget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*providerBudget, *modelBudget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") // Should fail at provider level (checked first) assertDecision(t, DecisionBudgetExceeded, result) assert.Contains(t, result.Reason, "Provider-level budget exceeded") } func TestResolver_EvaluateModelAndProviderRequest_ProviderRateLimitThenModelRateLimit(t *testing.T) { logger := NewMockLogger() // Provider rate limit exceeded providerRateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max provider := buildProviderWithGovernance("openai", nil, providerRateLimit) // Model rate limit within limit modelRateLimit := buildRateLimit("rl2", 20000, 2000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, modelRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*providerRateLimit, *modelRateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") // Should fail at provider level (checked first) assertDecision(t, DecisionTokenLimited, result) assert.Contains(t, result.Reason, "Provider-level rate limit check failed") } func TestResolver_EvaluateModelAndProviderRequest_ProviderRateLimitThenModelRateLimit_RequestLimit(t *testing.T) { logger := NewMockLogger() // Provider rate limit exceeded (request limit) providerRateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max provider := buildProviderWithGovernance("openai", nil, providerRateLimit) // Model rate limit within limit modelRateLimit := buildRateLimit("rl2", 20000, 2000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, modelRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*providerRateLimit, *modelRateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") // Should fail at provider level (checked first) assertDecision(t, DecisionRequestLimited, result) assert.Contains(t, result.Reason, "Provider-level rate limit check failed") } func TestResolver_EvaluateModelAndProviderRequest_AllChecksPass(t *testing.T) { logger := NewMockLogger() // Provider budget and rate limit within limits providerBudget := buildBudget("budget1", 100.0, "1h") providerRateLimit := buildRateLimit("rl1", 10000, 1000) provider := buildProviderWithGovernance("openai", providerBudget, providerRateLimit) // Model budget and rate limit within limits modelBudget := buildBudget("budget2", 200.0, "1h") modelRateLimit := buildRateLimit("rl2", 20000, 2000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, modelBudget, modelRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*providerBudget, *modelBudget}, RateLimits: []configstoreTables.TableRateLimit{*providerRateLimit, *modelRateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "gpt-4") assertDecision(t, DecisionAllow, result) assert.Contains(t, result.Reason, "provider-level and model-level checks passed") } func TestResolver_EvaluateModelAndProviderRequest_ProviderOnly_NoModel(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // No model provided result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.OpenAI, "") assertDecision(t, DecisionAllow, result) } func TestResolver_EvaluateModelAndProviderRequest_ModelOnly_NoProvider(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // No provider provided result := resolver.EvaluateModelAndProviderRequest(ctx, "", "gpt-4") assertDecision(t, DecisionAllow, result) } func TestResolver_EvaluateModelAndProviderRequest_ProviderSpecificBudget_DifferentProvider_Passes(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has budget (exceeded) budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // Request with Azure (different provider) for same model should pass result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.Azure, "gpt-4o") assertDecision(t, DecisionAllow, result) } func TestResolver_EvaluateModelAndProviderRequest_ProviderSpecificRateLimit_DifferentProvider_Passes(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // Request with Azure (different provider) for same model should pass result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.Azure, "gpt-4o") assertDecision(t, DecisionAllow, result) } func TestResolver_EvaluateModelAndProviderRequest_ProviderSpecificRateLimit_DifferentProvider_Passes_RequestLimit(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (request limit exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) resolver := NewBudgetResolver(store, nil, logger, nil) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // Request with Azure (different provider) for same model should pass result := resolver.EvaluateModelAndProviderRequest(ctx, schemas.Azure, "gpt-4o") assertDecision(t, DecisionAllow, result) } // ============================================================================ // End-to-End Tests - PreLLMHook Integration // ============================================================================ func TestPreLLMHook_ProviderBudgetExceeded_NoVirtualKey(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when provider budget is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "budget exceeded") } func TestPreLLMHook_ProviderRateLimitExceeded_NoVirtualKey(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when provider rate limit is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "rate limit") } func TestPreLLMHook_ModelBudgetExceeded_NoVirtualKey(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model budget is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "budget exceeded") } func TestPreLLMHook_ModelRateLimitExceeded_NoVirtualKey(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model rate limit is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "rate limit") } func TestPreLLMHook_ModelRateLimitExceeded_NoVirtualKey_RequestLimit(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model rate limit (request limit) is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "rate limit") } func TestPreLLMHook_AllChecksPass_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Provider budget and rate limit within limits providerBudget := buildBudget("budget1", 100.0, "1h") providerRateLimit := buildRateLimit("rl1", 10000, 1000) provider := buildProviderWithGovernance("openai", providerBudget, providerRateLimit) // Model budget and rate limit within limits modelBudget := buildBudget("budget2", 200.0, "1h") modelRateLimit := buildRateLimit("rl2", 20000, 2000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, modelBudget, modelRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*providerBudget, *modelBudget}, RateLimits: []configstoreTables.TableRateLimit{*providerRateLimit, *modelRateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } result, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.Nil(t, shortCircuit, "Should not short circuit when all checks pass") assert.NotNil(t, result) } func TestPreLLMHook_ProviderBudgetThenModelBudget_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Provider budget exceeded providerBudget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") provider := buildProviderWithGovernance("openai", providerBudget, nil) // Model budget within limit modelBudget := buildBudget("budget2", 200.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, modelBudget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*providerBudget, *modelBudget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) // Should fail at provider level (checked first) assert.NotNil(t, shortCircuit, "Should short circuit when provider budget is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "budget exceeded") } func TestPreLLMHook_ProviderSpecificModelBudget_DifferentProvider_Passes_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has budget (exceeded) budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.Azure, // Different provider Model: "gpt-4o", // Same model }, } result, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.Nil(t, shortCircuit, "Should not short circuit when model config is provider-specific and different provider is used") assert.NotNil(t, result) } func TestPreLLMHook_ProviderSpecificModelRateLimit_DifferentProvider_Passes_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Tokens at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.Azure, // Different provider Model: "gpt-4o", // Same model }, } result, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.Nil(t, shortCircuit, "Should not short circuit when model config is provider-specific and different provider is used") assert.NotNil(t, result) } func TestPreLLMHook_ProviderSpecificModelRateLimit_DifferentProvider_Passes_NoVirtualKey_RequestLimit(t *testing.T) { logger := NewMockLogger() // OpenAI GPT-4O has rate limit (request limit exceeded) rateLimit := buildRateLimitWithUsage("rl1", 10000, 0, 1000, 1000) // Requests at max providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.Azure, // Different provider Model: "gpt-4o", // Same model }, } result, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.Nil(t, shortCircuit, "Should not short circuit when model config is provider-specific and different provider is used (request limit)") assert.NotNil(t, result) } // ============================================================================ // End-to-End Tests - PreLLMHook Integration with Virtual Key Fallback // ============================================================================ func TestPreLLMHook_ModelProviderPass_VirtualKeyBudgetExceeded(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key budget exceeded vkBudget := buildBudgetWithUsage("vk-budget1", 100.0, 100.1, "1h") // Over limit vk := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", vkBudget) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, Budgets: []configstoreTables.TableBudget{*vkBudget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK budget is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "budget exceeded") } func TestPreLLMHook_ModelProviderPass_VirtualKeyRateLimitExceeded_Token(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key rate limit exceeded (token) vkRateLimit := buildRateLimitWithUsage("vk-rl1", 10000, 10000, 1000, 0) // Tokens at max vk := buildVirtualKeyWithRateLimit("vk1", "sk-bf-test", "Test VK", vkRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, RateLimits: []configstoreTables.TableRateLimit{*vkRateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK token rate limit is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "rate limit") } func TestPreLLMHook_ModelProviderPass_VirtualKeyRateLimitExceeded_Request(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key rate limit exceeded (request) vkRateLimit := buildRateLimitWithUsage("vk-rl1", 10000, 0, 1000, 1000) // Requests at max vk := buildVirtualKeyWithRateLimit("vk1", "sk-bf-test", "Test VK", vkRateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, RateLimits: []configstoreTables.TableRateLimit{*vkRateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK request rate limit is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "rate limit") } func TestPreLLMHook_ModelProviderPass_VirtualKeyChecksPass(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key checks also pass vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", true) vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{ buildProviderConfig("openai", []string{"*"}), } store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } result, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.Nil(t, shortCircuit, "Should not short circuit when both model/provider and VK checks pass") assert.NotNil(t, result) } func TestPreLLMHook_ModelProviderPass_VirtualKeyNotFound(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key not found store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-nonexistent") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK is not found") assert.Contains(t, shortCircuit.Error.Error.Message, "not found") } func TestPreLLMHook_ModelProviderPass_VirtualKeyBlocked(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key is inactive vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", false) // Inactive store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK is inactive") assert.Contains(t, shortCircuit.Error.Error.Message, "inactive") } func TestPreLLMHook_ModelProviderPass_VirtualKeyProviderBlocked(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key blocks OpenAI provider providerConfigs := []configstoreTables.TableVirtualKeyProviderConfig{ buildProviderConfig("anthropic", []string{"claude-3-sonnet"}), // Only Anthropic allowed } vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", providerConfigs) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, // Not allowed by VK Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK blocks provider") assert.Contains(t, shortCircuit.Error.Error.Message, "not allowed") } func TestPreLLMHook_ModelProviderPass_VirtualKeyModelBlocked(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (no limits) // Virtual key blocks specific model providerConfigs := []configstoreTables.TableVirtualKeyProviderConfig{ buildProviderConfig("openai", []string{"gpt-4", "gpt-4-turbo"}), // Only these models allowed } vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", providerConfigs) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-3.5-turbo", // Not in allowed list }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK blocks model") assert.Contains(t, shortCircuit.Error.Error.Message, "not allowed") } func TestPreLLMHook_ModelProviderPass_VirtualKeyBudgetExceeded_WithModelProviderLimits(t *testing.T) { logger := NewMockLogger() // Model/provider checks pass (within limits) providerBudget := buildBudget("provider-budget1", 200.0, "1h") provider := buildProviderWithGovernance("openai", providerBudget, nil) modelBudget := buildBudget("model-budget1", 150.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, modelBudget, nil) // Virtual key budget exceeded vkBudget := buildBudgetWithUsage("vk-budget1", 100.0, 100.1, "1h") // Over limit vk := buildVirtualKeyWithBudget("vk1", "sk-bf-test", "Test VK", vkBudget) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, Budgets: []configstoreTables.TableBudget{*providerBudget, *modelBudget, *vkBudget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) parentCtx := context.WithValue(context.Background(), schemas.BifrostContextKeyVirtualKey, "sk-bf-test") parentCtx = context.WithValue(parentCtx, schemas.BifrostContextKeyRequestID, "req-1") ctx := schemas.NewBifrostContext(parentCtx, schemas.NoDeadline) req := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit, _ := plugin.PreLLMHook(ctx, req) assert.NotNil(t, shortCircuit, "Should short circuit when model/provider pass but VK budget is exceeded") assert.Contains(t, shortCircuit.Error.Error.Message, "budget exceeded") } // ============================================================================ // End-to-End Tests - PostHook Integration (Usage Tracking) // ============================================================================ func TestPostHook_UpdatesProviderBudgetUsage_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Set budget with initial usage close to limit to test the flow // Note: Without model catalog, cost will be 0, so we test the flow even if budget isn't actually updated budget := buildBudgetWithUsage("budget1", 100.0, 50.0, "1h") provider := buildProviderWithGovernance("openai", budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) // First request: PreLLMHook should pass, PostHook updates usage parentCtx1 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-1") ctx1 := schemas.NewBifrostContext(parentCtx1, schemas.NoDeadline) req1 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit1, _ := plugin.PreLLMHook(ctx1, req1) assert.Nil(t, shortCircuit1, "First request should pass PreLLMHook") result1 := &schemas.BifrostResponse{ ChatResponse: &schemas.BifrostChatResponse{ Model: "gpt-4", Usage: &schemas.BifrostLLMUsage{ PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500, }, ExtraFields: schemas.BifrostResponseExtraFields{ RequestType: schemas.ChatCompletionRequest, Provider: schemas.OpenAI, OriginalModelRequested: "gpt-4", }, }, } _, _, err = plugin.PostLLMHook(ctx1, result1, nil) assert.NoError(t, err, "Should successfully process PostHook for provider budget usage update") // Wait for async processing to complete time.Sleep(200 * time.Millisecond) // Second request: Verify the flow works (budget check should still pass since cost is 0 without model catalog) parentCtx2 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-2") ctx2 := schemas.NewBifrostContext(parentCtx2, schemas.NoDeadline) req2 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit2, _ := plugin.PreLLMHook(ctx2, req2) // Without model catalog, cost is 0, so budget won't be exceeded // This test verifies the PostHook -> PreLLMHook flow works correctly assert.Nil(t, shortCircuit2, "Second request should pass PreLLMHook (cost is 0 without model catalog)") } func TestPostHook_UpdatesProviderRateLimitUsage_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Set rate limit: 10000 tokens, 1000 requests // First request: 10000 tokens, 1 request (brings usage to exactly the limit) // Second request: Should fail because we're already at the limit rateLimit := buildRateLimit("rl1", 10000, 1000) provider := buildProviderWithGovernance("openai", nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ Providers: []configstoreTables.TableProvider{*provider}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) // First request: PreLLMHook should pass, PostHook updates usage to 10000 parentCtx1 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-1") ctx1 := schemas.NewBifrostContext(parentCtx1, schemas.NoDeadline) req1 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit1, _ := plugin.PreLLMHook(ctx1, req1) assert.Nil(t, shortCircuit1, "First request should pass PreLLMHook") result1 := &schemas.BifrostResponse{ ChatResponse: &schemas.BifrostChatResponse{ Model: "gpt-4", Usage: &schemas.BifrostLLMUsage{ PromptTokens: 6000, CompletionTokens: 4000, TotalTokens: 10000, // 10000 tokens used (exactly at limit) }, ExtraFields: schemas.BifrostResponseExtraFields{ RequestType: schemas.ChatCompletionRequest, Provider: schemas.OpenAI, OriginalModelRequested: "gpt-4", }, }, } _, _, err = plugin.PostLLMHook(ctx1, result1, nil) assert.NoError(t, err, "Should successfully process PostHook for provider rate limit usage update") // Wait for async processing to complete time.Sleep(200 * time.Millisecond) // Second request: Should fail because we're already at the token limit (10000/10000) parentCtx2 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-2") ctx2 := schemas.NewBifrostContext(parentCtx2, schemas.NoDeadline) req2 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit2, _ := plugin.PreLLMHook(ctx2, req2) assert.NotNil(t, shortCircuit2, "Second request should fail PreLLMHook due to token limit exceeded") assert.Contains(t, shortCircuit2.Error.Error.Message, "token limit exceeded", "Error should indicate token limit exceeded") } func TestPostHook_UpdatesModelBudgetUsage_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Set budget with initial usage close to limit to test the flow // Note: Without model catalog, cost will be 0, so we test the flow even if budget isn't actually updated budget := buildBudgetWithUsage("budget1", 100.0, 50.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4", nil, budget, nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) // First request: PreLLMHook should pass, PostHook updates usage parentCtx1 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-1") ctx1 := schemas.NewBifrostContext(parentCtx1, schemas.NoDeadline) req1 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit1, _ := plugin.PreLLMHook(ctx1, req1) assert.Nil(t, shortCircuit1, "First request should pass PreLLMHook") result1 := &schemas.BifrostResponse{ ChatResponse: &schemas.BifrostChatResponse{ Model: "gpt-4", Usage: &schemas.BifrostLLMUsage{ PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500, }, ExtraFields: schemas.BifrostResponseExtraFields{ RequestType: schemas.ChatCompletionRequest, Provider: schemas.OpenAI, OriginalModelRequested: "gpt-4", }, }, } _, _, err = plugin.PostLLMHook(ctx1, result1, nil) assert.NoError(t, err, "Should successfully process PostHook for model budget usage update") // Wait for async processing to complete time.Sleep(200 * time.Millisecond) // Second request: Verify the flow works (budget check should still pass since cost is 0 without model catalog) parentCtx2 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-2") ctx2 := schemas.NewBifrostContext(parentCtx2, schemas.NoDeadline) req2 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit2, _ := plugin.PreLLMHook(ctx2, req2) // Without model catalog, cost is 0, so budget won't be exceeded // This test verifies the PostHook -> PreLLMHook flow works correctly assert.Nil(t, shortCircuit2, "Second request should pass PreLLMHook (cost is 0 without model catalog)") } func TestPostHook_UpdatesModelRateLimitUsage_NoVirtualKey(t *testing.T) { logger := NewMockLogger() // Set rate limit: 10000 tokens, 1000 requests // First request: 10000 tokens, 1 request (brings usage to exactly the limit) // Second request: Should fail because we're already at the limit rateLimit := buildRateLimit("rl1", 10000, 1000) modelConfig := buildModelConfig("mc1", "gpt-4", nil, nil, rateLimit) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, nil) require.NoError(t, err) plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) require.NoError(t, err) // First request: PreLLMHook should pass, PostHook updates usage to 10000 parentCtx1 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-1") ctx1 := schemas.NewBifrostContext(parentCtx1, schemas.NoDeadline) req1 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit1, _ := plugin.PreLLMHook(ctx1, req1) assert.Nil(t, shortCircuit1, "First request should pass PreLLMHook") result1 := &schemas.BifrostResponse{ ChatResponse: &schemas.BifrostChatResponse{ Model: "gpt-4", Usage: &schemas.BifrostLLMUsage{ PromptTokens: 6000, CompletionTokens: 4000, TotalTokens: 10000, // 10000 tokens used (exactly at limit) }, ExtraFields: schemas.BifrostResponseExtraFields{ RequestType: schemas.ChatCompletionRequest, Provider: schemas.OpenAI, OriginalModelRequested: "gpt-4", }, }, } _, _, err = plugin.PostLLMHook(ctx1, result1, nil) assert.NoError(t, err, "Should successfully process PostHook for model rate limit usage update") // Wait for async processing to complete time.Sleep(200 * time.Millisecond) // Second request: Should fail because we're already at the token limit (10000/10000) parentCtx2 := context.WithValue(context.Background(), schemas.BifrostContextKeyRequestID, "req-2") ctx2 := schemas.NewBifrostContext(parentCtx2, schemas.NoDeadline) req2 := &schemas.BifrostRequest{ RequestType: schemas.ChatCompletionRequest, ChatRequest: &schemas.BifrostChatRequest{ Provider: schemas.OpenAI, Model: "gpt-4", }, } _, shortCircuit2, _ := plugin.PreLLMHook(ctx2, req2) assert.NotNil(t, shortCircuit2, "Second request should fail PreLLMHook due to token limit exceeded") assert.Contains(t, shortCircuit2.Error.Error.Message, "token limit exceeded", "Error should indicate token limit exceeded") } // ============================================================================ // Cross-Provider Model Matching Tests // ============================================================================ // TestStore_CheckModelBudget_CrossProviderModelMatch tests that a model-only config // for "gpt-4o" is matched when the request uses "openai/gpt-4o" (OpenRouter-style prefix). func TestStore_CheckModelBudget_CrossProviderModelMatch(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, mc) require.NoError(t, err) // Request with provider-prefixed model name should match the "gpt-4o" config _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil) assert.Error(t, err, "Should reject: openai/gpt-4o should match model-only config for gpt-4o") assert.Contains(t, err.Error(), "budget exceeded") } // TestStore_CheckModelBudget_CrossProviderModelMatch_WithinLimit tests that the match works // and correctly allows requests within the budget. func TestStore_CheckModelBudget_CrossProviderModelMatch_WithinLimit(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, mc) require.NoError(t, err) _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil) assert.NoError(t, err, "Should allow: budget is within limit") } // TestStore_CheckModelRateLimit_CrossProviderModelMatch tests that a model-only rate limit config // for "gpt-4o" is matched when the request uses "openai/gpt-4o". func TestStore_CheckModelRateLimit_CrossProviderModelMatch(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Token limit at max modelConfig := buildModelConfig("mc1", "gpt-4o", nil, nil, rateLimit) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, mc) require.NoError(t, err) decision, errResult := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil, nil) assert.Error(t, errResult, "Should reject: openai/gpt-4o should match model-only rate limit for gpt-4o") assert.Contains(t, errResult.Error(), "token limit exceeded") assert.NotEqual(t, DecisionAllow, decision) } // TestStore_UpdateModelBudgetUsage_CrossProviderModelMatch tests that usage for "openai/gpt-4o" // is correctly attributed to the model-only config for "gpt-4o". func TestStore_UpdateModelBudgetUsage_CrossProviderModelMatch(t *testing.T) { logger := NewMockLogger() budget := buildBudget("budget1", 100.0, "1h") modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, mc) require.NoError(t, err) // Update usage with prefixed model name err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "openai/gpt-4o", schemas.OpenRouter, 50.0) assert.NoError(t, err, "Should successfully update budget usage via cross-provider match") // Now exceed the budget err = store.UpdateProviderAndModelBudgetUsageInMemory(context.Background(), "openai/gpt-4o", schemas.OpenRouter, 55.0) assert.NoError(t, err) // Budget should now be exceeded _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil) assert.Error(t, err, "Budget should be exceeded after usage updates via cross-provider match") assert.Contains(t, err.Error(), "budget exceeded") } // TestStore_UpdateModelRateLimitUsage_CrossProviderModelMatch tests that rate limit usage // for "openai/gpt-4o" is correctly attributed to the model-only config for "gpt-4o". func TestStore_UpdateModelRateLimitUsage_CrossProviderModelMatch(t *testing.T) { logger := NewMockLogger() rateLimit := buildRateLimitWithUsage("rl1", 100, 0, 1000, 0) // Low token limit modelConfig := buildModelConfig("mc1", "gpt-4o", nil, nil, rateLimit) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, }, mc) require.NoError(t, err) // Update token usage with prefixed model name err = store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "openai/gpt-4o", schemas.OpenRouter, 100, true, false) assert.NoError(t, err, "Should successfully update rate limit via cross-provider match") // Rate limit should now be exceeded decision, errResult := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil, nil) assert.Error(t, errResult, "Token limit should be exceeded after usage update via cross-provider match") assert.Contains(t, errResult.Error(), "token limit exceeded") assert.NotEqual(t, DecisionAllow, decision) } // TestStore_CheckModelBudget_ModelWithProvider_ExactMatchOnly tests that model+provider configs // (e.g., "gpt-4o:openai") use exact matching and do NOT fuzzy-match. func TestStore_CheckModelBudget_ModelWithProvider_ExactMatchOnly(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit providerStr := "openai" modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, budget, nil) mc := newTestModelCatalog(t) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, mc) require.NoError(t, err) // Request with the exact matching model+provider should be rejected (budget exceeded) _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenAI}, nil) assert.Error(t, err, "Exact model+provider match should apply budget") // Request with a different provider should NOT match the provider-specific config _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenRouter}, nil) assert.NoError(t, err, "Different provider should not match provider-specific config") } // TestStore_CheckModelBudget_NoCatalog_NoMatch tests that without a model catalog, // cross-provider matching does not happen (graceful degradation). func TestStore_CheckModelBudget_NoCatalog_NoMatch(t *testing.T) { logger := NewMockLogger() budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil) // No model catalog passed (nil) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig}, Budgets: []configstoreTables.TableBudget{*budget}, }, nil) require.NoError(t, err) // Without catalog, "openai/gpt-4o" won't match "gpt-4o" config _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "openai/gpt-4o", Provider: schemas.OpenRouter}, nil) assert.NoError(t, err, "Without model catalog, cross-provider matching should not happen") // Direct match should still work _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenAI}, nil) assert.Error(t, err, "Direct match should still work without catalog") }