mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-16 15:43:13 +08:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
432a6a461f | ||
|
|
c2d5bbd93e | ||
|
|
ae501cc220 |
@@ -187,16 +187,17 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t
|
|||||||
result.BalanceOverdrafted = !sufficient
|
result.BalanceOverdrafted = !sufficient
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Key 已不存在时跳过其自身的额度/限速计数,其余结算项不受影响。
|
||||||
if cmd.APIKeyQuotaCost > 0 {
|
if cmd.APIKeyQuotaCost > 0 {
|
||||||
exhausted, err := incrementUsageBillingAPIKeyQuota(ctx, tx, cmd.APIKeyID, cmd.APIKeyQuotaCost)
|
exhausted, err := incrementUsageBillingAPIKeyQuota(ctx, tx, cmd.APIKeyID, cmd.APIKeyQuotaCost)
|
||||||
if err != nil {
|
if err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
result.APIKeyQuotaExhausted = exhausted
|
result.APIKeyQuotaExhausted = exhausted
|
||||||
}
|
}
|
||||||
|
|
||||||
if cmd.APIKeyRateLimitCost > 0 {
|
if cmd.APIKeyRateLimitCost > 0 {
|
||||||
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil {
|
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -162,6 +162,61 @@ func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) {
|
|||||||
require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
|
require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestUsageBillingRepositoryApply_DeletedAPIKeyStillBillsBalance(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
client := testEntClient(t)
|
||||||
|
repo := NewUsageBillingRepository(client, integrationDB)
|
||||||
|
|
||||||
|
user := mustCreateUser(t, client, &service.User{
|
||||||
|
Email: fmt.Sprintf("usage-billing-deleted-key-user-%d@example.com", time.Now().UnixNano()),
|
||||||
|
PasswordHash: "hash",
|
||||||
|
Balance: 100,
|
||||||
|
})
|
||||||
|
apiKey := mustCreateApiKey(t, client, &service.APIKey{
|
||||||
|
UserID: user.ID,
|
||||||
|
Key: "sk-usage-billing-deleted-key-" + uuid.NewString(),
|
||||||
|
Name: "billing-deleted-key",
|
||||||
|
Quota: 50,
|
||||||
|
RateLimit5h: 50,
|
||||||
|
})
|
||||||
|
account := mustCreateAccount(t, client, &service.Account{
|
||||||
|
Name: "usage-billing-deleted-key-account-" + uuid.NewString(),
|
||||||
|
Type: service.AccountTypeAPIKey,
|
||||||
|
})
|
||||||
|
|
||||||
|
_, err := integrationDB.ExecContext(ctx, "UPDATE api_keys SET deleted_at = NOW() WHERE id = $1", apiKey.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
requestID := uuid.NewString()
|
||||||
|
result, err := repo.Apply(ctx, &service.UsageBillingCommand{
|
||||||
|
RequestID: requestID,
|
||||||
|
APIKeyID: apiKey.ID,
|
||||||
|
UserID: user.ID,
|
||||||
|
AccountID: account.ID,
|
||||||
|
AccountType: service.AccountTypeAPIKey,
|
||||||
|
BalanceCost: 1.25,
|
||||||
|
APIKeyQuotaCost: 1.25,
|
||||||
|
APIKeyRateLimitCost: 1.25,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, result)
|
||||||
|
require.True(t, result.Applied)
|
||||||
|
require.False(t, result.APIKeyQuotaExhausted)
|
||||||
|
|
||||||
|
var balance float64
|
||||||
|
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
|
||||||
|
require.InDelta(t, 98.75, balance, 0.000001)
|
||||||
|
|
||||||
|
var quotaUsed, usage5h float64
|
||||||
|
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used, usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan("aUsed, &usage5h))
|
||||||
|
require.InDelta(t, 0, quotaUsed, 0.000001)
|
||||||
|
require.InDelta(t, 0, usage5h, 0.000001)
|
||||||
|
|
||||||
|
var dedupCount int
|
||||||
|
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount))
|
||||||
|
require.Equal(t, 1, dedupCount)
|
||||||
|
}
|
||||||
|
|
||||||
func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
|
func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
client := testEntClient(t)
|
client := testEntClient(t)
|
||||||
|
|||||||
@@ -20,6 +20,8 @@ const (
|
|||||||
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`
|
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`
|
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`
|
userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL`
|
||||||
|
apiKeyQuotaIncrementSQL = `(?s)UPDATE api_keys\s+SET quota_used = quota_used \+ \$1,.*WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING`
|
||||||
|
apiKeyRateLimitIncrementSQL = `(?s)UPDATE api_keys SET\s+usage_5h = .*WHERE id = \$2 AND deleted_at IS NULL`
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) {
|
func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) {
|
||||||
@@ -99,6 +101,66 @@ func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) {
|
|||||||
require.NoError(t, mock.ExpectationsWereMet())
|
require.NoError(t, mock.ExpectationsWereMet())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestApplyUsageBillingEffects_DeletedAPIKeyStillBillsBalance(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)).
|
||||||
|
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(90.0))
|
||||||
|
mock.ExpectQuery(apiKeyQuotaIncrementSQL).
|
||||||
|
WithArgs(10.0, int64(7), service.StatusAPIKeyActive, service.StatusAPIKeyQuotaExhausted).
|
||||||
|
WillReturnError(sql.ErrNoRows)
|
||||||
|
mock.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||||
|
WithArgs(10.0, int64(7)).
|
||||||
|
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||||
|
mock.ExpectCommit()
|
||||||
|
|
||||||
|
result := &service.UsageBillingApplyResult{Applied: true}
|
||||||
|
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||||
|
UserID: 42,
|
||||||
|
APIKeyID: 7,
|
||||||
|
BalanceCost: 10,
|
||||||
|
APIKeyQuotaCost: 10,
|
||||||
|
APIKeyRateLimitCost: 10,
|
||||||
|
}, result)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, result.NewBalance)
|
||||||
|
require.InDelta(t, 90.0, *result.NewBalance, 0.000001)
|
||||||
|
require.False(t, result.APIKeyQuotaExhausted)
|
||||||
|
require.NoError(t, tx.Commit())
|
||||||
|
require.NoError(t, mock.ExpectationsWereMet())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyUsageBillingEffects_APIKeyCounterErrorStillFails(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.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||||
|
WithArgs(10.0, int64(7)).
|
||||||
|
WillReturnError(sql.ErrConnDone)
|
||||||
|
mock.ExpectRollback()
|
||||||
|
|
||||||
|
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||||
|
UserID: 42,
|
||||||
|
APIKeyID: 7,
|
||||||
|
APIKeyRateLimitCost: 10,
|
||||||
|
}, &service.UsageBillingApplyResult{Applied: true})
|
||||||
|
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||||
|
require.NoError(t, tx.Rollback())
|
||||||
|
require.NoError(t, mock.ExpectationsWereMet())
|
||||||
|
}
|
||||||
|
|
||||||
func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) {
|
func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) {
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
db, mock, err := sqlmock.New()
|
db, mock, err := sqlmock.New()
|
||||||
|
|||||||
Reference in New Issue
Block a user