feat(billing): 完成异步结算与请求执行闭环
统一任务成功、Attempt 与结算 Outbox 的事务边界，增加多实例安全结算、人工复核、请求幂等与执行租约。钱包决策使用九位精确金额，并通过审计保护约束保留流水事实；同时补充管理接口、指标与 PostgreSQL 集成测试。
This commit is contained in:
@@ -0,0 +1,265 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
func TestTaskIdempotencyAndExecutionLease(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-v2-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
input := CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "simulation",
|
||||
Request: map[string]any{"model": "billing-v2-model", "prompt": "lease"},
|
||||
IdempotencyKeyHash: "key-hash-" + uuid.NewString(), IdempotencyRequestHash: "request-a",
|
||||
}
|
||||
created, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || created.Replayed {
|
||||
t.Fatalf("create task result=%+v err=%v", created, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.Task.ID)
|
||||
})
|
||||
|
||||
replayed, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || !replayed.Replayed || replayed.Task.ID != created.Task.ID {
|
||||
t.Fatalf("replay result=%+v err=%v", replayed, err)
|
||||
}
|
||||
input.IdempotencyRequestHash = "request-b"
|
||||
if _, err := db.CreateTaskIdempotent(ctx, input, user); !errors.Is(err, ErrIdempotencyKeyReused) {
|
||||
t.Fatalf("different request error=%v", err)
|
||||
}
|
||||
|
||||
firstToken := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.Task.ID, firstToken, 5*time.Minute)
|
||||
if err != nil || claimed.ExecutionToken != firstToken {
|
||||
t.Fatalf("first claim task=%+v err=%v", claimed, err)
|
||||
}
|
||||
if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, firstToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("renew first lease: %v", err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.Task.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.Task.ID, secondToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("take over expired lease: %v", err)
|
||||
}
|
||||
if _, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: firstToken, Code: "old_worker", Message: "old"}); !errors.Is(err, ErrTaskExecutionLeaseLost) {
|
||||
t.Fatalf("old worker terminal error=%v", err)
|
||||
}
|
||||
finished, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: secondToken, Code: "new_worker", Message: "new"})
|
||||
if err != nil || finished.ErrorCode != "new_worker" {
|
||||
t.Fatalf("new worker terminal task=%+v err=%v", finished, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingSettlementStaleTakeoverDebitsExactlyOnce(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-settle-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{GatewayUserID: gatewayUserID, Currency: "resource", Balance: 10, Reason: "billing v2 test"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{Kind: "images.generations", Model: "billing-v2-model", RunMode: "production", Request: map[string]any{"model": "billing-v2-model"}}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
token := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
amount := "1.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, claimed, user, []any{map[string]any{"currency": "resource", "amount": amount}}, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource", "requestFingerprint": "billing-v2-test",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.FinishTaskSuccess(ctx, FinishTaskSuccessInput{
|
||||
TaskID: created.ID, ExecutionToken: token, Result: map[string]any{"ok": true},
|
||||
FinalChargeAmountText: amount, BillingCurrency: "resource",
|
||||
PricingSnapshot: map[string]any{"pricingVersion": "effective-pricing-v2"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
firstClaims, err := db.ClaimBillingSettlements(ctx, "worker-one", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first := settlementForTask(t, firstClaims, created.ID)
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE settlement_outbox SET locked_at=now()-interval '3 minutes' WHERE id=$1::uuid`, first.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondClaims, err := db.ClaimBillingSettlements(ctx, "worker-two", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second := settlementForTask(t, secondClaims, created.ID)
|
||||
if err := db.ProcessBillingSettlement(ctx, first); !errors.Is(err, pgx.ErrNoRows) {
|
||||
t.Fatalf("stale settlement lock error=%v", err)
|
||||
}
|
||||
if err := db.ProcessBillingSettlement(ctx, second); err != nil {
|
||||
t.Fatalf("process takeover: %v", err)
|
||||
}
|
||||
|
||||
var walletExact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT balance = 8.999999999::numeric
|
||||
AND frozen_balance = 0::numeric
|
||||
AND total_spent = 1.000000001::numeric
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID).Scan(&walletExact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !walletExact {
|
||||
t.Fatal("wallet amounts did not preserve nine-decimal settlement")
|
||||
}
|
||||
var billingTransactions int
|
||||
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_wallet_transactions WHERE reference_id=$1 AND transaction_type='task_billing'`, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions != 1 {
|
||||
t.Fatalf("task billing transactions=%d", billingTransactions)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `
|
||||
DELETE FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID); err == nil {
|
||||
t.Fatal("wallet account with audit transactions must not be deletable")
|
||||
}
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT count(*) FROM gateway_wallet_transactions
|
||||
WHERE gateway_user_id=$1::uuid AND reference_id=$2`, gatewayUserID, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions == 0 {
|
||||
t.Fatal("wallet audit transactions were lost after rejected account deletion")
|
||||
}
|
||||
settled, err := db.GetTask(ctx, created.ID)
|
||||
if err != nil || settled.BillingStatus != "settled" {
|
||||
t.Fatalf("settled task=%+v err=%v", settled, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseTaskBillingReservationsPreservesNineDecimalPlaces(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-release-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", Balance: 1, Reason: "billing v2 release test",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
amount := "0.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, created, user, nil, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0", Reason: "must not cross frozen balance",
|
||||
}); !errors.Is(err, ErrBalanceBelowFrozen) {
|
||||
t.Fatalf("balance below frozen error=%v", err)
|
||||
}
|
||||
if err := db.ReleaseTaskBillingReservations(ctx, reservations, "integration_test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0.123456789", Reason: "exact adjustment",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.RechargeUserWalletBalance(ctx, WalletRechargeInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", AmountText: "0.000000001", Reason: "exact recharge",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var exact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT account.balance = 0.123456790::numeric
|
||||
AND account.frozen_balance = 0::numeric
|
||||
AND task.reservation_amount = 0::numeric
|
||||
AND task.billing_status = 'not_started'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions transaction
|
||||
WHERE transaction.reference_id = task.id::text
|
||||
AND transaction.transaction_type = 'release'
|
||||
AND transaction.amount = 0.000000001::numeric
|
||||
)
|
||||
FROM gateway_wallet_accounts account
|
||||
JOIN gateway_tasks task ON task.gateway_user_id = account.gateway_user_id
|
||||
WHERE task.id=$1::uuid AND account.currency='resource'`, created.ID).Scan(&exact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !exact {
|
||||
t.Fatal("nine-decimal reservation was not released exactly")
|
||||
}
|
||||
}
|
||||
|
||||
func billingV2IntegrationStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run billing v2 PostgreSQL integration tests")
|
||||
}
|
||||
db, err := Connect(context.Background(), databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var databaseName string
|
||||
if err := db.pool.QueryRow(context.Background(), `SELECT current_database()`).Scan(&databaseName); err != nil {
|
||||
db.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
||||
db.Close()
|
||||
t.Fatalf("refusing to use non-test database %q", databaseName)
|
||||
}
|
||||
t.Cleanup(db.Close)
|
||||
return db
|
||||
}
|
||||
|
||||
func settlementForTask(t *testing.T, items []BillingSettlement, taskID string) BillingSettlement {
|
||||
t.Helper()
|
||||
for _, item := range items {
|
||||
if item.TaskID == taskID {
|
||||
return item
|
||||
}
|
||||
}
|
||||
t.Fatalf("settlement for task %s not claimed", taskID)
|
||||
return BillingSettlement{}
|
||||
}
|
||||
Reference in New Issue
Block a user