李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
262 lines
10 KiB
Go
262 lines
10 KiB
Go
//go:build unit
|
|
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"testing"
|
|
|
|
"github.com/DATA-DOG/go-sqlmock"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
)
|
|
|
|
const (
|
|
conditionalBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND balance >= \$1\s+RETURNING balance`
|
|
overdraftBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING balance`
|
|
reserveBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) \+ \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND balance >= \$1\s+RETURNING balance, frozen_balance`
|
|
captureBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance\s+\+ CASE WHEN \$1 > \$2 THEN \$1 - \$2 ELSE 0 END\s+- CASE WHEN \$2 > \$1 THEN \$2 - \$1 ELSE 0 END,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$3 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
|
releaseBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance \+ \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
|
userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL`
|
|
)
|
|
|
|
func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
|
WithArgs(2.5, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(7.5))
|
|
mock.ExpectCommit()
|
|
|
|
newBalance, sufficient, err := deductUsageBillingBalance(ctx, tx, 42, 2.5)
|
|
require.NoError(t, err)
|
|
require.True(t, sufficient)
|
|
require.InDelta(t, 7.5, newBalance, 0.000001)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestDeductUsageBillingBalance_RecordsOverdraftWhenGuardMisses(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectQuery(overdraftBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(-5.0))
|
|
mock.ExpectCommit()
|
|
|
|
newBalance, sufficient, err := deductUsageBillingBalance(ctx, tx, 42, 10)
|
|
require.NoError(t, err)
|
|
require.False(t, sufficient)
|
|
require.InDelta(t, -5.0, newBalance, 0.000001)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectQuery(overdraftBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(-5.0))
|
|
mock.ExpectCommit()
|
|
|
|
result := &service.UsageBillingApplyResult{Applied: true}
|
|
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
|
UserID: 42,
|
|
BalanceCost: 10,
|
|
}, result)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result.NewBalance)
|
|
require.InDelta(t, -5.0, *result.NewBalance, 0.000001)
|
|
require.True(t, result.BalanceOverdrafted)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectQuery(overdraftBalanceDeductSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectRollback()
|
|
|
|
_, _, err = deductUsageBillingBalance(ctx, tx, 42, 10)
|
|
require.ErrorIs(t, err, service.ErrUserNotFound)
|
|
require.NoError(t, tx.Rollback())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestReserveUsageBillingBatchImageBalance_MovesAvailableToFrozen(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(reserveBatchImageHoldSQL).
|
|
WithArgs(2.5, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(7.5, 2.5))
|
|
mock.ExpectCommit()
|
|
|
|
result, err := reserveUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 2.5})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result.NewBalance)
|
|
require.NotNil(t, result.FrozenBalance)
|
|
require.InDelta(t, 7.5, *result.NewBalance, 0.000001)
|
|
require.InDelta(t, 2.5, *result.FrozenBalance, 0.000001)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestReserveUsageBillingBatchImageBalance_InsufficientBalance(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(reserveBatchImageHoldSQL).
|
|
WithArgs(10.0, int64(42)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectQuery(userExistsForBillingSQL).
|
|
WithArgs(int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"?column?"}).AddRow(1))
|
|
mock.ExpectRollback()
|
|
|
|
_, err = reserveUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 10})
|
|
require.ErrorIs(t, err, service.ErrBatchImageInsufficientBalance)
|
|
require.NoError(t, tx.Rollback())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestCaptureUsageBillingBatchImageBalance_ReleasesRemainder(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(captureBatchImageHoldSQL).
|
|
WithArgs(1.0, 0.25, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(9.75, 0.0))
|
|
mock.ExpectCommit()
|
|
|
|
result, err := captureUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 1, ActualAmount: 0.25})
|
|
require.NoError(t, err)
|
|
require.InDelta(t, 9.75, *result.NewBalance, 0.000001)
|
|
require.InDelta(t, 0.0, *result.FrozenBalance, 0.000001)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestCaptureUsageBillingBatchImageBalance_RejectsActualCostOverHold(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectRollback()
|
|
|
|
_, err = captureUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 0.5, ActualAmount: 1})
|
|
require.ErrorIs(t, err, service.ErrBatchImageSettlementCostExceedsHold)
|
|
require.NoError(t, tx.Rollback())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestReleaseUsageBillingBatchImageBalance_ReturnsFrozenToAvailable(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`).
|
|
WithArgs(service.BatchImageHoldRequestID("imgbatch_release"), int64(7)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"?column?"}).AddRow(1))
|
|
mock.ExpectQuery(releaseBatchImageHoldSQL).
|
|
WithArgs(1.0, int64(42)).
|
|
WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(10.0, 0.0))
|
|
mock.ExpectCommit()
|
|
|
|
result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_release", HoldAmount: 1})
|
|
require.NoError(t, err)
|
|
require.InDelta(t, 10.0, *result.NewBalance, 0.000001)
|
|
require.InDelta(t, 0.0, *result.FrozenBalance, 0.000001)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|
|
|
|
func TestReleaseUsageBillingBatchImageBalance_SkipsWhenHoldNeverReserved(t *testing.T) {
|
|
ctx := context.Background()
|
|
db, mock, err := sqlmock.New()
|
|
require.NoError(t, err)
|
|
defer func() { _ = db.Close() }()
|
|
|
|
mock.ExpectBegin()
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
// dedup 与归档表均无 hold claim：说明该 job 从未成功冻结，
|
|
// 释放必须跳过，不得从他人冻结资金池中凭空生成余额。
|
|
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`).
|
|
WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup_archive\s+WHERE request_id = \$1 AND api_key_id = \$2`).
|
|
WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)).
|
|
WillReturnError(sql.ErrNoRows)
|
|
mock.ExpectCommit()
|
|
|
|
result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_phantom", HoldAmount: 1})
|
|
require.NoError(t, err)
|
|
require.Nil(t, result.NewBalance)
|
|
require.Nil(t, result.FrozenBalance)
|
|
require.NoError(t, tx.Commit())
|
|
require.NoError(t, mock.ExpectationsWereMet())
|
|
}
|