summaryrefslogtreecommitdiff
path: root/internal/billing/service_test.go
diff options
context:
space:
mode:
authorChia <Chia@93.nz>2026-08-06 09:29:41 +1200
committerChia <Chia@93.nz>2026-08-06 09:32:46 +1200
commit41e322c53d7b4b796eb377d0df9c29ecd10ba431 (patch)
treec730526150e55e39b822d5197e4a20318ecaa449 /internal/billing/service_test.go
parenteadb2ffe85c43cf6fc741c9823cd28eedb4a844c (diff)
feat: complete commercial control plane, billing, auth, and model catalog
- add PostgreSQL control-plane persistence with Redis-degraded hot reload - implement prepaid balance, usage ledger, Stripe top-up and reconciliation - add registration, email verification, password reset, invitations and RBAC - support TOTP, Passkey MFA, device sessions, quotas and rate limits - add tenant billing profiles, audit logs and operational readiness checks - build authenticated admin console, Quickstart, Playground and usage analytics - add public model catalog with pricing, filtering and cost estimation - support OpenAI Responses providers and provider health failover - validate real upstream usage reporting and balance settlement
Diffstat (limited to 'internal/billing/service_test.go')
-rw-r--r--internal/billing/service_test.go67
1 files changed, 67 insertions, 0 deletions
diff --git a/internal/billing/service_test.go b/internal/billing/service_test.go
index 6a1bb49..5e21b20 100644
--- a/internal/billing/service_test.go
+++ b/internal/billing/service_test.go
@@ -4,6 +4,7 @@ import (
"context"
"crypto/sha256"
"encoding/json"
+ "errors"
"fmt"
"net/http"
"net/http/httptest"
@@ -19,6 +20,72 @@ import (
"github.com/stripe/stripe-go/v86/webhook"
)
+func TestAuthorizeEnforcesAPIKeyMonthlySpendCapPostgres(t *testing.T) {
+ databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL")
+ if databaseURL == "" {
+ t.Skip("AIGW_TEST_DATABASE_URL is not set")
+ }
+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+ defer cancel()
+ if err := controlplane.MigrateDatabase(ctx, databaseURL); err != nil {
+ t.Fatal(err)
+ }
+ service, err := New(ctx, Options{DatabaseURL: databaseURL, Currency: "usd", DefaultMaxOutputTokens: 10})
+ if err != nil {
+ t.Fatal(err)
+ }
+ t.Cleanup(service.Close)
+ slug := fmt.Sprintf("key-budget-%d", time.Now().UnixNano())
+ var tenantID, projectID, keyID string
+ if err := service.db.QueryRow(ctx, `INSERT INTO tenants (slug,name) VALUES ($1,'Key budget') RETURNING id::text`, slug).Scan(&tenantID); err != nil {
+ t.Fatal(err)
+ }
+ if err := service.db.QueryRow(ctx, `INSERT INTO projects (tenant_id,slug,name) VALUES ($1,'default','Default') RETURNING id::text`, tenantID).Scan(&projectID); err != nil {
+ t.Fatal(err)
+ }
+ hash := sha256.Sum256([]byte(slug))
+ if err := service.db.QueryRow(ctx, `INSERT INTO api_keys (tenant_id,project_id,name,key_prefix,key_hash,monthly_spend_micros) VALUES ($1,$2,'limited','sk-test',$3,9) RETURNING id::text`, tenantID, projectID, hash[:]).Scan(&keyID); err != nil {
+ t.Fatal(err)
+ }
+ t.Cleanup(func() {
+ for _, statement := range []struct {
+ query string
+ args []any
+ }{
+ {`DELETE FROM billing_settlement_jobs WHERE request_id LIKE 'req_key_budget_%'`, nil},
+ {`DELETE FROM billing_reservations WHERE tenant_id=$1`, []any{tenantID}},
+ {`DELETE FROM tenant_wallets WHERE tenant_id=$1`, []any{tenantID}},
+ {`DELETE FROM api_keys WHERE id=$1`, []any{keyID}},
+ {`DELETE FROM projects WHERE id=$1`, []any{projectID}},
+ {`DELETE FROM tenants WHERE id=$1`, []any{tenantID}},
+ } {
+ if _, cleanupErr := service.db.Exec(context.Background(), statement.query, statement.args...); cleanupErr != nil {
+ t.Errorf("cleanup API key budget test data: %v", cleanupErr)
+ }
+ }
+ })
+ if _, err := service.db.Exec(ctx, `INSERT INTO tenant_wallets (tenant_id,currency,balance_micros) VALUES ($1,'usd',1000000)`, tenantID); err != nil {
+ t.Fatal(err)
+ }
+ model := domain.Model{ID: "model/budget", PriceCurrency: "usd", OutputPriceMicrosPerMillion: 1_000_000}
+ principal := domain.Principal{TenantID: tenantID, ProjectID: projectID, KeyID: keyID, MonthlySpendMicros: 9}
+ err = service.Authorize(ctx, Authorization{RequestID: "req_key_budget_rejected", Principal: principal, Model: model, Body: []byte(`{"max_tokens":10}`)})
+ if !errors.Is(err, ErrQuotaExceeded) {
+ t.Fatalf("Authorize error = %v, want ErrQuotaExceeded", err)
+ }
+ var rejectedReservations int
+ if err := service.db.QueryRow(ctx, `SELECT count(*) FROM billing_reservations WHERE request_id='req_key_budget_rejected'`).Scan(&rejectedReservations); err != nil {
+ t.Fatal(err)
+ }
+ if rejectedReservations != 0 {
+ t.Fatalf("quota rejection left %d reservation rows", rejectedReservations)
+ }
+ principal.MonthlySpendMicros = 10
+ if err := service.Authorize(ctx, Authorization{RequestID: "req_key_budget_allowed", Principal: principal, Model: model, Body: []byte(`{"max_tokens":10}`)}); err != nil {
+ t.Fatalf("Authorize at exact cap: %v", err)
+ }
+}
+
func TestUsageCostUsesFixedPointAndRoundsOnce(t *testing.T) {
usage := domain.Usage{InputTokens: 3, OutputTokens: 2, CacheReadInputTokens: 5}
cost, err := usageCost(usage, 150_000, 600_000, 30_000, 0)