417 lines
15 KiB
Go
417 lines
15 KiB
Go
package billing
|
||
|
||
import (
|
||
"context"
|
||
"os"
|
||
"strconv"
|
||
"sync"
|
||
"testing"
|
||
|
||
"github.com/alicebob/miniredis/v2"
|
||
"github.com/gogf/gf/v2/database/gredis"
|
||
"github.com/shopspring/decimal"
|
||
|
||
// 注册 GoFrame redis 适配器（生产代码在 main.go 空导入；单测需在此显式注册）
|
||
_ "github.com/gogf/gf/contrib/nosql/redis/v2"
|
||
)
|
||
|
||
// miniredis 全局实例：g.Redis() 经 gredis.SetConfig 指向内存 Redis，
|
||
// 验证钱包权威化架构下的全部 Lua 脚本（预扣/结算认领/解冻/加款/扣款）。
|
||
// miniredis 单线程串行执行命令，脚本内 check-then-act 等价原子，可验证资金守恒。
|
||
var mr *miniredis.Miniredis
|
||
|
||
func TestMain(m *testing.M) {
|
||
var err error
|
||
mr, err = miniredis.Run()
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
if err := gredis.SetConfigByMap(map[string]any{
|
||
"address": mr.Addr(),
|
||
"db": 0,
|
||
}); err != nil {
|
||
panic(err)
|
||
}
|
||
code := m.Run()
|
||
mr.Close()
|
||
os.Exit(code)
|
||
}
|
||
|
||
func testCtx() context.Context { return context.Background() }
|
||
|
||
// readWallet 读取租户钱包余额/冻结（micro → decimal）
|
||
func readWallet(t *testing.T, tenantID int64) (balance, frozen decimal.Decimal) {
|
||
t.Helper()
|
||
b, f, _, _, exists, err := readWalletHash(testCtx(), tenantID)
|
||
if err != nil {
|
||
t.Fatalf("read wallet hash: %v", err)
|
||
}
|
||
if !exists {
|
||
t.Fatalf("wallet hash missing for tenant %d", tenantID)
|
||
}
|
||
return FromMicro(b), FromMicro(f)
|
||
}
|
||
|
||
// readWalletConsumed 读取租户累计消费计数（micro → decimal；hash 缺字段按 0）
|
||
func readWalletConsumed(t *testing.T, tenantID int64) decimal.Decimal {
|
||
t.Helper()
|
||
_, _, c, hasConsumed, exists, err := readWalletHash(testCtx(), tenantID)
|
||
if err != nil {
|
||
t.Fatalf("read wallet hash: %v", err)
|
||
}
|
||
if !exists {
|
||
t.Fatalf("wallet hash missing for tenant %d", tenantID)
|
||
}
|
||
if !hasConsumed {
|
||
return decimal.Zero
|
||
}
|
||
return FromMicro(c)
|
||
}
|
||
|
||
// seedWallet 用加款创建指定余额的钱包 hash（顺带校验 CreditWalletRedis 自身）
|
||
func seedWallet(t *testing.T, tenantID int64, balance decimal.Decimal) {
|
||
t.Helper()
|
||
ba, _, err := CreditWalletRedis(testCtx(), tenantID, balance)
|
||
if err != nil {
|
||
t.Fatalf("seed wallet: %v", err)
|
||
}
|
||
if !ba.Equal(balance) {
|
||
t.Fatalf("seed wallet balance=%s want %s", ba.String(), balance.String())
|
||
}
|
||
}
|
||
|
||
// assertDecimalEqual 断言两个 decimal 相等（金额精度 6 位）
|
||
func assertDecimalEqual(t *testing.T, got, want decimal.Decimal, msg string) {
|
||
t.Helper()
|
||
if !got.Equal(want) {
|
||
t.Fatalf("%s: got=%s want=%s", msg, got.String(), want.String())
|
||
}
|
||
}
|
||
|
||
func TestPreDeduct_Basic(t *testing.T) {
|
||
tenant := int64(1)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
|
||
ok, err := PreDeduct(testCtx(), tenant, 4, "req-p1", "gpt-4o")
|
||
if err != nil || !ok {
|
||
t.Fatalf("pre-deduct failed: ok=%v err=%v", ok, err)
|
||
}
|
||
bal, frozen := readWallet(t, tenant)
|
||
assertDecimalEqual(t, bal, NewFromFloat(10), "balance unchanged by freeze")
|
||
assertDecimalEqual(t, frozen, NewFromFloat(4), "frozen after prededuct")
|
||
|
||
// 幂等：同一 request_id 重复预扣 → 成功且不重复冻结
|
||
ok, err = PreDeduct(testCtx(), tenant, 4, "req-p1", "gpt-4o")
|
||
if err != nil || !ok {
|
||
t.Fatalf("idempotent pre-deduct failed: ok=%v err=%v", ok, err)
|
||
}
|
||
_, frozen = readWallet(t, tenant)
|
||
assertDecimalEqual(t, frozen, NewFromFloat(4), "no double freeze on idempotent")
|
||
|
||
// 余额不足：可用 6（10-4）< 7 → 拒绝且不冻结
|
||
ok, err = PreDeduct(testCtx(), tenant, 7, "req-p2", "gpt-4o")
|
||
if ok || err == nil {
|
||
t.Fatalf("insufficient pre-deduct should fail: ok=%v err=%v", ok, err)
|
||
}
|
||
_, frozen = readWallet(t, tenant)
|
||
assertDecimalEqual(t, frozen, NewFromFloat(4), "no freeze on insufficient")
|
||
}
|
||
|
||
func TestSettleClaim_Normal(t *testing.T) {
|
||
tenant := int64(2)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 4, "req-s1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
|
||
claimed, bal, frozen, err := SettleClaim(testCtx(), tenant, NewFromFloat(1.5), []string{"req-s1"})
|
||
if err != nil {
|
||
t.Fatalf("settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, claimed, NewFromFloat(4), "claimed == prededucted")
|
||
assertDecimalEqual(t, bal, NewFromFloat(8.5), "balance after deduct")
|
||
assertDecimalEqual(t, frozen, decimal.Zero, "frozen fully released")
|
||
}
|
||
|
||
func TestSettleClaim_Duplicate(t *testing.T) {
|
||
tenant := int64(3)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 4, "req-s2", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
if _, _, _, err := SettleClaim(testCtx(), tenant, NewFromFloat(1.5), []string{"req-s2"}); err != nil {
|
||
t.Fatalf("first settle: %v", err)
|
||
}
|
||
|
||
// 重复结算：预扣 hash 已随认领删除 → claimed=0，只按补扣继续扣 balance
|
||
claimed, bal, frozen, err := SettleClaim(testCtx(), tenant, NewFromFloat(0.5), []string{"req-s2"})
|
||
if err != nil {
|
||
t.Fatalf("duplicate settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, claimed, decimal.Zero, "duplicate claims nothing")
|
||
assertDecimalEqual(t, bal, NewFromFloat(8), "only supplement deducted")
|
||
assertDecimalEqual(t, frozen, decimal.Zero, "frozen untouched")
|
||
}
|
||
|
||
func TestSettleClaim_MultiKeyAdjust(t *testing.T) {
|
||
tenant := int64(4)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 3, "req-t1", "midjourney"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 1, "req-t1_adjust", "midjourney"); !ok {
|
||
t.Fatal("adjust prededuct failed")
|
||
}
|
||
|
||
claimed, bal, frozen, err := SettleClaim(testCtx(), tenant, NewFromFloat(2), []string{"req-t1", "req-t1_adjust"})
|
||
if err != nil {
|
||
t.Fatalf("multi-key settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, claimed, NewFromFloat(4), "sum of both prededucts claimed")
|
||
assertDecimalEqual(t, bal, NewFromFloat(8), "balance after cost")
|
||
assertDecimalEqual(t, frozen, decimal.Zero, "all frozen released")
|
||
}
|
||
|
||
func TestUnfreeze_Idempotent(t *testing.T) {
|
||
tenant := int64(5)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 2, "req-u1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
|
||
amt, err := UnfreezePreDeduct(testCtx(), tenant, "req-u1")
|
||
if err != nil {
|
||
t.Fatalf("unfreeze: %v", err)
|
||
}
|
||
assertDecimalEqual(t, amt, NewFromFloat(2), "unfreeze releases prededucted amount")
|
||
_, frozen := readWallet(t, tenant)
|
||
assertDecimalEqual(t, frozen, decimal.Zero, "frozen released")
|
||
|
||
// 幂等：第二次认领不到 → 0，不二次释放
|
||
amt2, err := UnfreezePreDeduct(testCtx(), tenant, "req-u1")
|
||
if err != nil {
|
||
t.Fatalf("second unfreeze: %v", err)
|
||
}
|
||
assertDecimalEqual(t, amt2, decimal.Zero, "idempotent second unfreeze")
|
||
|
||
// 不存在的预扣 → 0
|
||
amt3, err := UnfreezePreDeduct(testCtx(), tenant, "req-missing")
|
||
if err != nil {
|
||
t.Fatalf("missing unfreeze: %v", err)
|
||
}
|
||
assertDecimalEqual(t, amt3, decimal.Zero, "missing prededuct no-op")
|
||
}
|
||
|
||
func TestCreditDebitWallet(t *testing.T) {
|
||
tenant := int64(6)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
|
||
// 扣款门槛：可用 10 < 12 → 拒绝且不变更
|
||
_, _, ok, err := DebitWalletRedis(testCtx(), tenant, NewFromFloat(12))
|
||
if err != nil {
|
||
t.Fatalf("debit err: %v", err)
|
||
}
|
||
if ok {
|
||
t.Fatal("debit beyond available must fail")
|
||
}
|
||
bal, _ := readWallet(t, tenant)
|
||
assertDecimalEqual(t, bal, NewFromFloat(10), "no change on failed debit")
|
||
|
||
// 正常扣款
|
||
after, _, ok, err := DebitWalletRedis(testCtx(), tenant, NewFromFloat(3))
|
||
if err != nil || !ok {
|
||
t.Fatalf("debit failed: ok=%v err=%v", ok, err)
|
||
}
|
||
assertDecimalEqual(t, after, NewFromFloat(7), "debit reduces balance")
|
||
|
||
// 加款（负数=扣回补偿）
|
||
after, _, err = CreditWalletRedis(testCtx(), tenant, NewFromFloat(-3))
|
||
if err != nil {
|
||
t.Fatalf("credit: %v", err)
|
||
}
|
||
assertDecimalEqual(t, after, NewFromFloat(4), "negative credit reduces balance")
|
||
}
|
||
|
||
// TestWalletConcurrency 并发守恒：N 并发预扣+结算后，frozen 归零、balance 精确扣减
|
||
func TestWalletConcurrency(t *testing.T) {
|
||
tenant := int64(10)
|
||
seedWallet(t, tenant, NewFromFloat(100))
|
||
|
||
const n = 50
|
||
var wg sync.WaitGroup
|
||
errCh := make(chan error, n)
|
||
for i := 0; i < n; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
ctx := testCtx()
|
||
rid := "req-c" + strconv.Itoa(idx)
|
||
ok, err := PreDeduct(ctx, tenant, 1, rid, "gpt-4o")
|
||
if err != nil || !ok {
|
||
errCh <- err
|
||
return
|
||
}
|
||
if _, _, _, err := SettleClaim(ctx, tenant, NewFromFloat(0.5), []string{rid}); err != nil {
|
||
errCh <- err
|
||
}
|
||
}(i)
|
||
}
|
||
wg.Wait()
|
||
close(errCh)
|
||
for err := range errCh {
|
||
t.Fatalf("concurrent wallet op failed: %v", err)
|
||
}
|
||
|
||
bal, frozen := readWallet(t, tenant)
|
||
assertDecimalEqual(t, bal, NewFromFloat(100).Sub(NewFromFloat(0.5).Mul(NewFromFloat(n))), "balance conserved (100 - n*0.5)")
|
||
assertDecimalEqual(t, frozen, decimal.Zero, "frozen converges to zero")
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(0.5).Mul(NewFromFloat(n)), "total_consumed conserved (n*0.5)")
|
||
}
|
||
|
||
// TestPredeductSweep_RecomputesFrozen 孤儿清扫重算：某预扣 hash 丢失（TTL 过期/崩溃残留）后，
|
||
// frozen 仍含其金额 → sweep 重算为幸存 hash 之和并清理活跃集合死亡成员。
|
||
func TestPredeductSweep_RecomputesFrozen(t *testing.T) {
|
||
tenant := int64(8)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 2, "req-sw1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct 1 failed")
|
||
}
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 3, "req-sw2", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct 2 failed")
|
||
}
|
||
_, frozen := readWallet(t, tenant)
|
||
assertDecimalEqual(t, frozen, NewFromFloat(5), "frozen after two prededucts")
|
||
|
||
// 模拟孤儿：直接删除 req-sw2 的预扣 hash（不释放冻结）
|
||
if !mr.Del(PreDeductRedisKeyPrefix + "req-sw2") {
|
||
t.Fatal("simulate orphan: prededuct hash not found")
|
||
}
|
||
|
||
PredeductSweep(testCtx())
|
||
|
||
_, frozen = readWallet(t, tenant)
|
||
assertDecimalEqual(t, frozen, NewFromFloat(2), "frozen recomputed to surviving prededucts")
|
||
|
||
// 孤儿 request 移出活跃集合
|
||
activeSetKey := "prededuct_active:8"
|
||
members, err := mr.SMembers(activeSetKey)
|
||
if err != nil {
|
||
t.Fatalf("read active set: %v", err)
|
||
}
|
||
if len(members) != 1 || members[0] != "req-sw1" {
|
||
t.Fatalf("active set after sweep = %v, want [req-sw1]", members)
|
||
}
|
||
}
|
||
|
||
func TestCleanupPreDeduct_RemovesHash(t *testing.T) {
|
||
tenant := int64(7)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 1, "req-cl1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
CleanupPreDeduct(testCtx(), tenant, "req-cl1")
|
||
if _, exists := GetPreDeductAmount(testCtx(), "req-cl1"); exists {
|
||
t.Fatal("prededuct hash should be cleaned")
|
||
}
|
||
}
|
||
|
||
// TestSettleClaim_IncrementsTotalConsumed 累计消费随结算扣款同一原子递增；
|
||
// 重复结算（预扣已认领、claimed=0 只补扣）时计数仍按实际扣款额累计。
|
||
func TestSettleClaim_IncrementsTotalConsumed(t *testing.T) {
|
||
tenant := int64(11)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 4, "req-tc1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
|
||
// 首次结算：新 hash 无 total_consumed 字段，HINCRBY 从 0 起算（正常路径由三层补种保证含基线）
|
||
if _, _, _, err := SettleClaim(testCtx(), tenant, NewFromFloat(1.5), []string{"req-tc1"}); err != nil {
|
||
t.Fatalf("settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(1.5), "total_consumed after first settle")
|
||
|
||
// 重复结算走补扣路径：计数继续累计实际扣款
|
||
if _, _, _, err := SettleClaim(testCtx(), tenant, NewFromFloat(0.5), []string{"req-tc1"}); err != nil {
|
||
t.Fatalf("second settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(2), "total_consumed accumulates supplement")
|
||
}
|
||
|
||
// TestReverseConsumeRedis_ReversesBalanceAndCount 结算补偿逆转：余额与累计消费同一原子回退
|
||
func TestReverseConsumeRedis_ReversesBalanceAndCount(t *testing.T) {
|
||
tenant := int64(12)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 4, "req-tc2", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
if _, _, _, err := SettleClaim(testCtx(), tenant, NewFromFloat(3), []string{"req-tc2"}); err != nil {
|
||
t.Fatalf("settle claim: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(3), "total_consumed after settle")
|
||
|
||
// 模拟 executeSettlement 第 c 步失败补偿：本单免费，计数与余额同步回退
|
||
if err := ReverseConsumeRedis(testCtx(), tenant, NewFromFloat(3)); err != nil {
|
||
t.Fatalf("reverse consume: %v", err)
|
||
}
|
||
bal, _ := readWallet(t, tenant)
|
||
assertDecimalEqual(t, bal, NewFromFloat(10), "balance restored after reverse")
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), decimal.Zero, "total_consumed reversed with balance")
|
||
}
|
||
|
||
// TestNonConsumeOps_DoNotTouchTotalConsumed 回归保护：
|
||
// 预扣/解冻/扣款（退款调账）/加款（充值）均不得触碰累计消费计数。
|
||
func TestNonConsumeOps_DoNotTouchTotalConsumed(t *testing.T) {
|
||
tenant := int64(13)
|
||
seedWallet(t, tenant, NewFromFloat(10))
|
||
|
||
if ok, _ := PreDeduct(testCtx(), tenant, 2, "req-nc1", "gpt-4o"); !ok {
|
||
t.Fatal("prededuct failed")
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), decimal.Zero, "prededuct must not touch total_consumed")
|
||
|
||
if _, err := UnfreezePreDeduct(testCtx(), tenant, "req-nc1"); err != nil {
|
||
t.Fatalf("unfreeze: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), decimal.Zero, "unfreeze must not touch total_consumed")
|
||
|
||
if _, _, ok, err := DebitWalletRedis(testCtx(), tenant, NewFromFloat(3)); err != nil || !ok {
|
||
t.Fatalf("debit failed: ok=%v err=%v", ok, err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), decimal.Zero, "debit (refund/adjust) must not touch total_consumed")
|
||
|
||
if _, _, err := CreditWalletRedis(testCtx(), tenant, NewFromFloat(5)); err != nil {
|
||
t.Fatalf("credit: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), decimal.Zero, "credit (recharge) must not touch total_consumed")
|
||
}
|
||
|
||
// TestSeedTotalConsumedIfMissing 三层补种之①③层共用 Lua 的行为：
|
||
// hash 存在且缺字段 → 写入基线；字段已存在 → 不覆盖（幂等）；hash 不存在 → no-op。
|
||
func TestSeedTotalConsumedIfMissing(t *testing.T) {
|
||
tenant := int64(14)
|
||
seedWallet(t, tenant, NewFromFloat(10)) // creditLua 创建的 hash 不含 total_consumed 字段
|
||
|
||
if err := seedTotalConsumedIfMissing(testCtx(), tenant, ToMicro(NewFromFloat(7.25))); err != nil {
|
||
t.Fatalf("seed total_consumed: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(7.25), "seeded from db baseline")
|
||
|
||
// 字段已存在后发生结算 → 重复种子不得回退结算增量
|
||
if _, _, _, err := SettleClaim(testCtx(), tenant, NewFromFloat(1), nil); err != nil {
|
||
t.Fatalf("settle claim: %v", err)
|
||
}
|
||
if err := seedTotalConsumedIfMissing(testCtx(), tenant, ToMicro(NewFromFloat(7.25))); err != nil {
|
||
t.Fatalf("re-seed: %v", err)
|
||
}
|
||
assertDecimalEqual(t, readWalletConsumed(t, tenant), NewFromFloat(8.25), "re-seed must not overwrite post-settle value")
|
||
|
||
// hash 不存在的租户 → no-op（不得凭空创建 hash，由 rebuild 流程负责）
|
||
missing := int64(15)
|
||
if err := seedTotalConsumedIfMissing(testCtx(), missing, ToMicro(NewFromFloat(9))); err != nil {
|
||
t.Fatalf("seed on missing hash: %v", err)
|
||
}
|
||
if _, _, _, hasConsumed, exists, err := readWalletHash(testCtx(), missing); err != nil || exists || hasConsumed {
|
||
t.Fatalf("seed must not create hash: exists=%v hasConsumed=%v err=%v", exists, hasConsumed, err)
|
||
}
|
||
}
|