1 files changed, 190 insertions, 0 deletions
diff --git a/internal/billing/service_test.go b/internal/billing/service_test.go
new file mode 100644
index 0000000..17f8143
--- /dev/null
+++ b/internal/billing/service_test.go
@@ -0,0 +1,190 @@
+package billing
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"aigw/internal/controlplane"
+	"aigw/internal/domain"
+
+	"github.com/stripe/stripe-go/v86"
+	"github.com/stripe/stripe-go/v86/webhook"
+)
+
+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)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// (3*150000 + 2*600000 + 5*30000) / 1e6 = 1.8 micros.
+	if cost != 2 {
+		t.Fatalf("cost = %d, want 2", cost)
+	}
+}
+
+func TestReservationUsesExplicitOutputLimit(t *testing.T) {
+	model := domain.Model{InputPriceMicrosPerMillion: 1_000_000, OutputPriceMicrosPerMillion: 1_000_000}
+	body := []byte(`{"max_tokens":8192}`)
+	cost, err := reservationCost(model, body, 4096)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := int64(len(body)) + 8192
+	if cost != want {
+		t.Fatalf("reservation = %d, want %d", cost, want)
+	}
+}
+
+func TestMinorToMicrosSupportsCurrencyExponents(t *testing.T) {
+	tests := []struct {
+		currency string
+		minor    int64
+		want     int64
+	}{{"usd", 123, 1_230_000}, {"jpy", 123, 123_000_000}, {"bhd", 123, 123_000}}
+	for _, test := range tests {
+		got, err := minorToMicros(test.currency, test.minor)
+		if err != nil {
+			t.Fatalf("%s: %v", test.currency, err)
+		}
+		if got != test.want {
+			t.Fatalf("%s: got %d, want %d", test.currency, got, test.want)
+		}
+	}
+}
+
+func TestCollectibleChargePreservesOtherReservations(t *testing.T) {
+	charged, err := collectibleCharge(100, 100, 80, 40)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if charged != 60 {
+		t.Fatalf("charged = %d, want 60", charged)
+	}
+	if remainingBalance, remainingHeld := int64(100)-charged, int64(80)-40; remainingBalance < remainingHeld {
+		t.Fatalf("remaining balance %d does not cover held %d", remainingBalance, remainingHeld)
+	}
+}
+
+func TestIntegrationIdentifierSuffixUsesLetters(t *testing.T) {
+	value := randomLetters(8)
+	if len(value) != 8 {
+		t.Fatalf("length = %d", len(value))
+	}
+	for _, char := range value {
+		if char < 'a' || char > 'z' {
+			t.Fatalf("non-letter suffix %q", value)
+		}
+	}
+}
+
+func TestWebhookRejectsInvalidSignatureBeforeProcessing(t *testing.T) {
+	service := &Service{stripeWebhookSecret: "whsec_test"}
+	request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(`{"id":"evt_fake"}`))
+	request.Header.Set("Stripe-Signature", "invalid")
+	response := httptest.NewRecorder()
+	service.WebhookHandler().ServeHTTP(response, request)
+	if response.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want 400", response.Code)
+	}
+}
+
+func TestStripeWebhookCreditsPaidOrderExactlyOncePostgres(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", MinTopUpMinor: 500, MaxTopUpMinor: 1_000_000,
+		StripeEnabled: true, StripeAPIKey: "rk_test_placeholder", StripeWebhookSecret: "whsec_integration_test",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(service.Close)
+
+	var tenantID string
+	slug := fmt.Sprintf("stripe-%d", time.Now().UnixNano())
+	if err := service.db.QueryRow(ctx, `INSERT INTO tenants (slug, name) VALUES ($1,'Stripe integration') RETURNING id::text`, slug).Scan(&tenantID); err != nil {
+		t.Fatal(err)
+	}
+	eventID := fmt.Sprintf("evt_aigw_%d", time.Now().UnixNano())
+	t.Cleanup(func() {
+		cleanupCtx := context.Background()
+		for _, statement := range []struct {
+			query string
+			arg   string
+		}{
+			{`DELETE FROM stripe_webhook_events WHERE event_id=$1`, eventID},
+			{`DELETE FROM billing_ledger WHERE tenant_id=$1`, tenantID},
+			{`DELETE FROM tenant_wallets WHERE tenant_id=$1`, tenantID},
+			{`DELETE FROM topup_orders WHERE tenant_id=$1`, tenantID},
+			{`DELETE FROM tenants WHERE id=$1`, tenantID},
+		} {
+			if _, cleanupErr := service.db.Exec(cleanupCtx, statement.query, statement.arg); cleanupErr != nil {
+				t.Errorf("cleanup Stripe integration data: %v", cleanupErr)
+			}
+		}
+	})
+	const amountMinor int64 = 2500
+	orderID, amountMicros, err := service.createTopUpOrder(ctx, CheckoutInput{TenantID: tenantID, AmountMinor: amountMinor})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sessionID := fmt.Sprintf("cs_test_aigw_%d", time.Now().UnixNano())
+	payload, err := json.Marshal(map[string]any{
+		"id": eventID, "object": "event", "api_version": stripe.APIVersion,
+		"type": string(stripe.EventTypeCheckoutSessionCompleted),
+		"data": map[string]any{"object": map[string]any{
+			"id": sessionID, "object": "checkout.session", "client_reference_id": orderID,
+			"amount_total": amountMinor, "currency": "usd", "payment_status": "paid",
+		}},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	signed := webhook.GenerateTestSignedPayload(&webhook.UnsignedPayload{Payload: payload, Secret: service.stripeWebhookSecret})
+	for delivery := 0; delivery < 2; delivery++ {
+		request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(string(payload)))
+		request.Header.Set("Stripe-Signature", signed.Header)
+		response := httptest.NewRecorder()
+		service.WebhookHandler().ServeHTTP(response, request)
+		if response.Code != http.StatusOK {
+			t.Fatalf("delivery %d status = %d, body = %s", delivery+1, response.Code, response.Body.String())
+		}
+	}
+
+	var balance int64
+	if err := service.db.QueryRow(ctx, `SELECT balance_micros FROM tenant_wallets WHERE tenant_id=$1`, tenantID).Scan(&balance); err != nil {
+		t.Fatal(err)
+	}
+	if balance != amountMicros {
+		t.Fatalf("balance = %d, want %d", balance, amountMicros)
+	}
+	var ledgerCount, webhookCount int
+	var orderStatus string
+	if err := service.db.QueryRow(ctx, `SELECT count(*) FROM billing_ledger WHERE source_type='stripe_checkout' AND source_id=$1`, sessionID).Scan(&ledgerCount); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.db.QueryRow(ctx, `SELECT count(*) FROM stripe_webhook_events WHERE event_id=$1`, eventID).Scan(&webhookCount); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.db.QueryRow(ctx, `SELECT status FROM topup_orders WHERE id=$1`, orderID).Scan(&orderStatus); err != nil {
+		t.Fatal(err)
+	}
+	if ledgerCount != 1 || webhookCount != 1 || orderStatus != "paid" {
+		t.Fatalf("ledger=%d webhook=%d order=%s, want 1/1/paid", ledgerCount, webhookCount, orderStatus)
+	}
+}