chengcheng
5b2b94b1bd
统一任务成功、Attempt 与结算 Outbox 的事务边界，增加多实例安全结算、人工复核、请求幂等与执行租约。钱包决策使用九位精确金额，并通过审计保护约束保留流水事实；同时补充管理接口、指标与 PostgreSQL 集成测试。
266 lines
11 KiB
Go
266 lines
11 KiB
Go
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{}
|
|
}
|