From 1a3d7f9a8a181df48f0e911cbe17a3fad3ab9ac9 Mon Sep 17 00:00:00 2001 From: Chia Date: Wed, 5 Aug 2026 00:26:25 +1200 Subject: add some scripts --- internal/billing/service_test.go | 190 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 internal/billing/service_test.go (limited to 'internal/billing/service_test.go') 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) + } +} -- cgit v1.2.3