package pg import ( "context" "sync" "testing" "time" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/store" ) func TestPGUsageCapStoreReserveUsageIdempotent(t *testing.T) { db := hooksTestDB(t) tenantID, _ := seedTenantAndAgent(t, db) usageStore := NewPGUsageCapStore(db) maxTokens := int64(100) policy := &store.UsageCapPolicy{ TenantID: tenantID, Window: store.UsageCapWindowHour, MaxTokens: &maxTokens, Enabled: true, Priority: 100, } if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err != nil { t.Fatalf("CreateUsageCapPolicy: %v", err) } req := store.UsageReserveRequest{ UsageCapScope: store.UsageCapScope{TenantID: tenantID}, ReservationKey: "duplicate-reservation", EstimatedTokens: 10, } for i := range 2 { if _, err := usageStore.ReserveUsage(context.Background(), req, []store.UsageCapPolicy{*policy}); err != nil { t.Fatalf("ReserveUsage call %d: %v", i+1, err) } } var reservedTokens, reservationRows int64 if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { t.Fatalf("query reserved tokens: %v", err) } if err := db.QueryRow(`SELECT COUNT(*) FROM usage_cap_reservations WHERE policy_id=$1 AND reservation_key=$2`, policy.ID, req.ReservationKey).Scan(&reservationRows); err != nil { t.Fatalf("query reservation rows: %v", err) } if reservedTokens != 10 { t.Fatalf("reserved_tokens = %d, want 10", reservedTokens) } if reservationRows != 1 { t.Fatalf("reservation rows = %d, want 1", reservationRows) } var wg sync.WaitGroup errs := make(chan error, 2) for range 2 { wg.Add(1) go func() { defer wg.Done() errs <- usageStore.ReconcileUsage(context.Background(), store.UsageReconcileRequest{ ReservationKey: req.ReservationKey, ActualTokens: 7, Status: "reconciled", }) }() } wg.Wait() close(errs) for err := range errs { if err != nil { t.Fatalf("ReconcileUsage: %v", err) } } var usedTokens int64 if err := db.QueryRow(`SELECT COALESCE(SUM(used_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&usedTokens); err != nil { t.Fatalf("query used tokens: %v", err) } if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { t.Fatalf("query reserved tokens after reconcile: %v", err) } if usedTokens != 7 { t.Fatalf("used_tokens = %d, want 7", usedTokens) } if reservedTokens != 0 { t.Fatalf("reserved_tokens after reconcile = %d, want 0", reservedTokens) } } func TestPGUsageCapStoreRejectsCrossTenantRefs(t *testing.T) { db := hooksTestDB(t) tenantA, agentA := seedTenantAndAgent(t, db) tenantB, _ := seedTenantAndAgent(t, db) usageStore := NewPGUsageCapStore(db) policy := &store.UsageCapPolicy{ TenantID: tenantB, AgentID: &agentA, Window: store.UsageCapWindowDay, MaxTokens: int64PtrPG(100), Enabled: true, } if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err == nil { t.Fatal("CreateUsageCapPolicy accepted agent_id from another tenant") } providerID := uuid.New() if _, err := db.Exec( `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) VALUES ($1,$2,$3,'openrouter','sk-test',true)`, providerID, tenantA, "ucp-"+providerID.String(), ); err != nil { t.Fatalf("seed provider: %v", err) } override := &store.UsagePricingOverride{ TenantID: tenantB, ProviderID: providerID, ProviderType: store.ProviderOpenRouter, ModelID: "openai/gpt-test", Enabled: true, } if err := usageStore.PutPricingOverride(context.Background(), override); err == nil { t.Fatal("PutPricingOverride accepted provider_id from another tenant") } masterProviderID := uuid.New() if _, err := db.Exec( `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) VALUES ($1,$2,$3,'openrouter','sk-test',true)`, masterProviderID, store.MasterTenantID, "ucpm-"+masterProviderID.String(), ); err != nil { t.Fatalf("seed master provider: %v", err) } t.Cleanup(func() { db.Exec("DELETE FROM llm_providers WHERE id=$1", masterProviderID) }) masterScopedPolicy := &store.UsageCapPolicy{ TenantID: tenantB, ProviderID: &masterProviderID, Window: store.UsageCapWindowDay, MaxTokens: int64PtrPG(100), Enabled: true, } if err := usageStore.CreateUsageCapPolicy(context.Background(), masterScopedPolicy); err != nil { t.Fatalf("CreateUsageCapPolicy rejected master provider ref: %v", err) } } func TestValidateUsagePricingFieldsRejectsNegativeValues(t *testing.T) { negative := "-0.01" if err := validateUsagePricingFields(store.UsagePricingFields{Input: &negative}); err == nil { t.Fatal("validateUsagePricingFields accepted negative price") } } func TestPGUsageCapStoreUpsertPricingCatalogSkipsInvalidEntries(t *testing.T) { db := hooksTestDB(t) usageStore := NewPGUsageCapStore(db) inputPrice := "0.000001" outputPrice := "0.000002" negativePrice := "-0.000001" entries := []store.UsagePricingCatalogEntry{ {ModelID: "invalid/negative-output", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &negativePrice}, SyncedAt: time.Now().UTC()}, {ModelID: "openai/gpt-valid", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, } count, err := usageStore.UpsertPricingCatalog(context.Background(), entries) if err != nil { t.Fatalf("UpsertPricingCatalog: %v", err) } if count != 1 { t.Fatalf("upserted count = %d, want 1", count) } listed, err := usageStore.ListPricingCatalog(context.Background(), store.UsagePricingQuery{ModelID: "gpt-valid"}) if err != nil { t.Fatalf("ListPricingCatalog valid: %v", err) } if len(listed) != 1 || listed[0].ModelID != "openai/gpt-valid" { t.Fatalf("valid entries = %+v, want openai/gpt-valid", listed) } listed, err = usageStore.ListPricingCatalog(context.Background(), store.UsagePricingQuery{ModelID: "negative-output"}) if err != nil { t.Fatalf("ListPricingCatalog invalid: %v", err) } if len(listed) != 0 { t.Fatalf("invalid entries = %+v, want none", listed) } } func TestPGUsageCapStoreResolvePricingUsesOpenRouterAliases(t *testing.T) { db := hooksTestDB(t) usageStore := NewPGUsageCapStore(db) inputPrice := "0.000001" outputPrice := "0.000002" entries := []store.UsagePricingCatalogEntry{ {ModelID: "openai/gpt-4o-mini", CanonicalModelID: "openai/gpt-4o-mini", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, {ModelID: "openai/gpt-5.5", CanonicalModelID: "openai/gpt-5.5", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, {ModelID: "anthropic/claude-3-5-haiku-latest", CanonicalModelID: "anthropic/claude-3-5-haiku-latest", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, {ModelID: "google/gemini-2.5-flash", CanonicalModelID: "google/gemini-2.5-flash", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, {ModelID: "qwen/qwen3.7-plus", CanonicalModelID: "qwen/qwen3.7-plus", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()}, } if _, err := usageStore.UpsertPricingCatalog(context.Background(), entries); err != nil { t.Fatalf("UpsertPricingCatalog: %v", err) } cases := []struct { name string providerName string providerType string modelID string wantModelID string }{ {name: "openai compat", providerName: "openai", providerType: store.ProviderOpenAICompat, modelID: "gpt-4o-mini", wantModelID: "openai/gpt-4o-mini"}, {name: "openai compat custom gpt model", providerName: "cppai", providerType: store.ProviderOpenAICompat, modelID: "gpt-5.5", wantModelID: "openai/gpt-5.5"}, {name: "anthropic native", providerName: "anthropic", providerType: store.ProviderAnthropicNative, modelID: "claude-3-5-haiku-latest", wantModelID: "anthropic/claude-3-5-haiku-latest"}, {name: "gemini native", providerName: "gemini", providerType: store.ProviderGeminiNative, modelID: "gemini-2.5-flash", wantModelID: "google/gemini-2.5-flash"}, {name: "bailian qwen model", providerName: "bailian", providerType: store.ProviderBailian, modelID: "qwen3.7-plus", wantModelID: "qwen/qwen3.7-plus"}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { resolved, err := usageStore.ResolvePricing(context.Background(), uuid.New(), uuid.New(), tc.providerName, tc.providerType, tc.modelID) if err != nil { t.Fatalf("ResolvePricing: %v", err) } if resolved.ModelID != tc.wantModelID { t.Fatalf("resolved model = %q, want %q", resolved.ModelID, tc.wantModelID) } }) } } func int64PtrPG(v int64) *int64 { return &v }