From 41e322c53d7b4b796eb377d0df9c29ecd10ba431 Mon Sep 17 00:00:00 2001
From: Chia <Chia@93.nz>
Date: Thu, 6 Aug 2026 09:29:41 +1200
Subject: 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
---
 internal/billing/service_test.go | 67 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

(limited to 'internal/billing/service_test.go')

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)
-- 
cgit v1.2.3