wans10andCaIon
58d4e9bd3b
* fix(billing): 异步任务退款时同步减少 used_quota 退款时仅恢复了 quota（剩余额度），但未同步减少 used_quota（已用额度）， 导致"总额度"（quota + used_quota）随退款次数持续虚增，超出用户实际充值金额。 修复三处退款路径： - RefundTaskQuota：任务失败完整退款 - RecalculateTaskQuota：差额结算退款分支 - controller/midjourney.go：Midjourney 任务失败退款 新增 model.UpdateUserUsedQuota 公开函数，仅调整 used_quota 不影响 request_count。 * fix(billing): 任务退款时同步扣减渠道 used_quota * fix(billing): complete async task refund accounting * style(model): group internal Midjourney fields --------- Co-authored-by: CaIon <i@caion.me>
1229 lines
40 KiB
Go
1229 lines
40 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"math"
|
|
"net/http"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/types"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/shopspring/decimal"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
if err != nil {
|
|
panic("failed to open test db: " + err.Error())
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
panic("failed to get sql.DB: " + err.Error())
|
|
}
|
|
sqlDB.SetMaxOpenConns(1)
|
|
|
|
model.DB = db
|
|
model.LOG_DB = db
|
|
|
|
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
|
common.RedisEnabled = false
|
|
common.BatchUpdateEnabled = false
|
|
common.LogConsumeEnabled = true
|
|
|
|
if err := db.AutoMigrate(
|
|
&model.Task{},
|
|
&model.User{},
|
|
&model.Token{},
|
|
&model.Log{},
|
|
&model.Channel{},
|
|
&model.Midjourney{},
|
|
&model.TopUp{},
|
|
&model.UserSubscription{},
|
|
&model.SystemTask{},
|
|
&model.SystemTaskLock{},
|
|
); err != nil {
|
|
panic("failed to migrate: " + err.Error())
|
|
}
|
|
|
|
os.Exit(m.Run())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Seed helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func truncate(t *testing.T) {
|
|
t.Helper()
|
|
t.Cleanup(func() {
|
|
model.DB.Exec("DELETE FROM tasks")
|
|
model.DB.Exec("DELETE FROM users")
|
|
model.DB.Exec("DELETE FROM tokens")
|
|
model.DB.Exec("DELETE FROM logs")
|
|
model.DB.Exec("DELETE FROM channels")
|
|
model.DB.Exec("DELETE FROM midjourneys")
|
|
model.DB.Exec("DELETE FROM top_ups")
|
|
model.DB.Exec("DELETE FROM user_subscriptions")
|
|
model.DB.Exec("DELETE FROM system_task_locks")
|
|
model.DB.Exec("DELETE FROM system_tasks")
|
|
})
|
|
}
|
|
|
|
func seedUser(t *testing.T, id int, quota int) {
|
|
t.Helper()
|
|
user := &model.User{Id: id, Username: "test_user", Quota: quota, Status: common.UserStatusEnabled}
|
|
require.NoError(t, model.DB.Create(user).Error)
|
|
}
|
|
|
|
func seedToken(t *testing.T, id int, userId int, key string, remainQuota int) {
|
|
t.Helper()
|
|
token := &model.Token{
|
|
Id: id,
|
|
UserId: userId,
|
|
Key: key,
|
|
Name: "test_token",
|
|
Status: common.TokenStatusEnabled,
|
|
RemainQuota: remainQuota,
|
|
UsedQuota: 0,
|
|
}
|
|
require.NoError(t, model.DB.Create(token).Error)
|
|
}
|
|
|
|
func seedSubscription(t *testing.T, id int, userId int, amountTotal int64, amountUsed int64) {
|
|
t.Helper()
|
|
sub := &model.UserSubscription{
|
|
Id: id,
|
|
UserId: userId,
|
|
AmountTotal: amountTotal,
|
|
AmountUsed: amountUsed,
|
|
Status: "active",
|
|
StartTime: time.Now().Unix(),
|
|
EndTime: time.Now().Add(30 * 24 * time.Hour).Unix(),
|
|
}
|
|
require.NoError(t, model.DB.Create(sub).Error)
|
|
}
|
|
|
|
func seedChannel(t *testing.T, id int) {
|
|
t.Helper()
|
|
ch := &model.Channel{Id: id, Name: "test_channel", Key: "sk-test", Status: common.ChannelStatusEnabled}
|
|
require.NoError(t, model.DB.Create(ch).Error)
|
|
}
|
|
|
|
func seedChargedAccounting(t *testing.T, userID, channelID, tokenID, quota, requestCount int) {
|
|
t.Helper()
|
|
require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Updates(map[string]any{
|
|
"used_quota": quota,
|
|
"request_count": requestCount,
|
|
}).Error)
|
|
require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", channelID).
|
|
Update("used_quota", quota).Error)
|
|
if tokenID > 0 {
|
|
require.NoError(t, model.DB.Model(&model.Token{}).Where("id = ?", tokenID).
|
|
Update("used_quota", quota).Error)
|
|
}
|
|
}
|
|
|
|
func makeTask(userId, channelId, quota, tokenId int, billingSource string, subscriptionId int) *model.Task {
|
|
return &model.Task{
|
|
TaskID: "task_" + time.Now().Format("150405.000"),
|
|
UserId: userId,
|
|
ChannelId: channelId,
|
|
Quota: quota,
|
|
Status: model.TaskStatus(model.TaskStatusInProgress),
|
|
Group: "default",
|
|
Data: json.RawMessage(`{}`),
|
|
CreatedAt: time.Now().Unix(),
|
|
UpdatedAt: time.Now().Unix(),
|
|
Properties: model.Properties{
|
|
OriginModelName: "test-model",
|
|
},
|
|
PrivateData: model.TaskPrivateData{
|
|
BillingSource: billingSource,
|
|
SubscriptionId: subscriptionId,
|
|
TokenId: tokenId,
|
|
BillingContext: &model.TaskBillingContext{
|
|
ModelPrice: 0.02,
|
|
GroupRatio: 1.0,
|
|
OriginModelName: "test-model",
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func TestPriceDataOtherRatiosFilterAndSnapshot(t *testing.T) {
|
|
priceData := types.PriceData{}
|
|
|
|
priceData.AddOtherRatio("zero", 0)
|
|
priceData.AddOtherRatio("negative", -0.5)
|
|
priceData.AddOtherRatio("nan", math.NaN())
|
|
priceData.AddOtherRatio("inf", math.Inf(1))
|
|
priceData.AddOtherRatio("one", 1)
|
|
priceData.AddOtherRatio("positive", 2.5)
|
|
|
|
ratios := priceData.OtherRatios()
|
|
require.Len(t, ratios, 2)
|
|
assert.Equal(t, 1.0, ratios["one"])
|
|
assert.Equal(t, 2.5, ratios["positive"])
|
|
assert.True(t, priceData.HasOtherRatio("one"))
|
|
assert.False(t, priceData.HasOtherRatio("zero"))
|
|
|
|
ratios["positive"] = 99
|
|
ratios["new"] = 3
|
|
nextSnapshot := priceData.OtherRatios()
|
|
assert.Equal(t, 2.5, nextSnapshot["positive"])
|
|
assert.NotContains(t, nextSnapshot, "new")
|
|
}
|
|
|
|
func TestPriceDataReplaceAndApplyOtherRatios(t *testing.T) {
|
|
priceData := types.PriceData{}
|
|
|
|
replaced := priceData.ReplaceOtherRatios(map[string]float64{
|
|
"zero": 0,
|
|
"negative": -3,
|
|
"nan": math.NaN(),
|
|
"inf": math.Inf(1),
|
|
"one": 1,
|
|
"duration": 2,
|
|
"size": 1.5,
|
|
})
|
|
|
|
require.True(t, replaced)
|
|
assert.Equal(t, 3.0, priceData.OtherRatioMultiplier())
|
|
assert.Equal(t, 30.0, priceData.ApplyOtherRatiosToFloat(10))
|
|
assert.Equal(t, 10.0, priceData.RemoveOtherRatiosFromFloat(30))
|
|
assert.True(t, decimal.NewFromInt(30).Equal(priceData.ApplyOtherRatiosToDecimal(decimal.NewFromInt(10))))
|
|
|
|
replaced = priceData.ReplaceOtherRatios(map[string]float64{
|
|
"zero": 0,
|
|
"nan": math.NaN(),
|
|
})
|
|
|
|
require.False(t, replaced)
|
|
assert.Nil(t, priceData.OtherRatios())
|
|
assert.Equal(t, 1.0, priceData.OtherRatioMultiplier())
|
|
}
|
|
|
|
func TestTaskBillingOtherFiltersHistoricalOtherRatios(t *testing.T) {
|
|
task := makeTask(1, 1, 100, 0, BillingSourceWallet, 0)
|
|
task.PrivateData.BillingContext.OtherRatios = map[string]float64{
|
|
"seconds": 2,
|
|
"identity": 1,
|
|
"zero": 0,
|
|
"negative": -1,
|
|
"nan": math.NaN(),
|
|
"inf": math.Inf(1),
|
|
}
|
|
|
|
other := taskBillingOther(task)
|
|
|
|
assert.Equal(t, 2.0, other["seconds"])
|
|
assert.Equal(t, 1.0, other["identity"])
|
|
assert.NotContains(t, other, "zero")
|
|
assert.NotContains(t, other, "negative")
|
|
assert.NotContains(t, other, "nan")
|
|
assert.NotContains(t, other, "inf")
|
|
}
|
|
|
|
func TestTaskBillingContextPriceDataFiltersMultiplier(t *testing.T) {
|
|
priceData := taskBillingContextPriceData(&model.TaskBillingContext{
|
|
OtherRatios: map[string]float64{
|
|
"seconds": 2,
|
|
"size": 3,
|
|
"identity": 1,
|
|
"zero": 0,
|
|
"negative": -1,
|
|
"nan": math.NaN(),
|
|
"inf": math.Inf(1),
|
|
},
|
|
})
|
|
|
|
require.NotNil(t, priceData)
|
|
assert.Equal(t, 6.0, priceData.OtherRatioMultiplier())
|
|
assert.Equal(t, map[string]float64{
|
|
"seconds": 2,
|
|
"size": 3,
|
|
"identity": 1,
|
|
}, priceData.OtherRatios())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Read-back helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func getUserQuota(t *testing.T, id int) int {
|
|
t.Helper()
|
|
var user model.User
|
|
require.NoError(t, model.DB.Select("quota").Where("id = ?", id).First(&user).Error)
|
|
return user.Quota
|
|
}
|
|
|
|
func getUserUsageAccounting(t *testing.T, id int) (int, int) {
|
|
t.Helper()
|
|
var user model.User
|
|
require.NoError(t, model.DB.Select("used_quota", "request_count").Where("id = ?", id).First(&user).Error)
|
|
return user.UsedQuota, user.RequestCount
|
|
}
|
|
|
|
func getChannelUsedQuota(t *testing.T, id int) int64 {
|
|
t.Helper()
|
|
var channel model.Channel
|
|
require.NoError(t, model.DB.Select("used_quota").Where("id = ?", id).First(&channel).Error)
|
|
return channel.UsedQuota
|
|
}
|
|
|
|
func getTokenRemainQuota(t *testing.T, id int) int {
|
|
t.Helper()
|
|
var token model.Token
|
|
require.NoError(t, model.DB.Select("remain_quota").Where("id = ?", id).First(&token).Error)
|
|
return token.RemainQuota
|
|
}
|
|
|
|
func getTokenUsedQuota(t *testing.T, id int) int {
|
|
t.Helper()
|
|
var token model.Token
|
|
require.NoError(t, model.DB.Select("used_quota").Where("id = ?", id).First(&token).Error)
|
|
return token.UsedQuota
|
|
}
|
|
|
|
func getSubscriptionUsed(t *testing.T, id int) int64 {
|
|
t.Helper()
|
|
var sub model.UserSubscription
|
|
require.NoError(t, model.DB.Select("amount_used").Where("id = ?", id).First(&sub).Error)
|
|
return sub.AmountUsed
|
|
}
|
|
|
|
func getTaskQuota(t *testing.T, id int64) int {
|
|
t.Helper()
|
|
var task model.Task
|
|
require.NoError(t, model.DB.Select("quota").Where("id = ?", id).First(&task).Error)
|
|
return task.Quota
|
|
}
|
|
|
|
func getMidjourneyTask(t *testing.T, id int) model.Midjourney {
|
|
t.Helper()
|
|
var task model.Midjourney
|
|
require.NoError(t, model.DB.First(&task, id).Error)
|
|
return task
|
|
}
|
|
|
|
func getLastLog(t *testing.T) *model.Log {
|
|
t.Helper()
|
|
var log model.Log
|
|
err := model.LOG_DB.Order("id desc").First(&log).Error
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
return &log
|
|
}
|
|
|
|
func countLogs(t *testing.T) int64 {
|
|
t.Helper()
|
|
var count int64
|
|
model.LOG_DB.Model(&model.Log{}).Count(&count)
|
|
return count
|
|
}
|
|
|
|
// ===========================================================================
|
|
// Legacy Midjourney billing tests
|
|
// ===========================================================================
|
|
|
|
func TestPrepareMidjourneyTaskBillingKeepsUnbilledMarkerClear(t *testing.T) {
|
|
task := &model.Midjourney{Quota: 900, TokenId: 7, BillingChannelId: 8}
|
|
|
|
prepared, err := PrepareMidjourneyTaskBilling(&relaycommon.RelayInfo{}, task, 900, false)
|
|
|
|
require.NoError(t, err)
|
|
assert.False(t, prepared)
|
|
assert.Zero(t, task.Quota)
|
|
assert.Zero(t, task.TokenId)
|
|
assert.Zero(t, task.BillingChannelId)
|
|
}
|
|
|
|
func TestSettleMidjourneyTaskBillingRequiresPersistedTask(t *testing.T) {
|
|
truncate(t)
|
|
|
|
const userID, tokenID, channelID = 49, 49, 49
|
|
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
|
seedUser(t, userID, initialUserQuota)
|
|
seedToken(t, tokenID, userID, "sk-midjourney-unpersisted", initialTokenQuota)
|
|
seedChannel(t, channelID)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
UserId: userID,
|
|
TokenId: tokenID,
|
|
TokenKey: "sk-midjourney-unpersisted",
|
|
UserQuota: initialUserQuota,
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelId: channelID,
|
|
},
|
|
}
|
|
task := &model.Midjourney{UserId: userID, ChannelId: channelID}
|
|
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
|
require.NoError(t, err)
|
|
require.True(t, prepared)
|
|
|
|
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
|
|
|
require.Error(t, err)
|
|
assert.False(t, billed)
|
|
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
|
}
|
|
|
|
func TestMidjourneyRefundRestoresEveryAccountingElementOnBillingChannel(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, billingChannelID, executionChannelID = 50, 50, 50, 51
|
|
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
|
seedUser(t, userID, initialUserQuota)
|
|
seedToken(t, tokenID, userID, "sk-midjourney", initialTokenQuota)
|
|
seedChannel(t, billingChannelID)
|
|
seedChannel(t, executionChannelID)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
UserId: userID,
|
|
TokenId: tokenID,
|
|
TokenKey: "sk-midjourney",
|
|
UserQuota: initialUserQuota,
|
|
UsingGroup: "default",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelId: billingChannelID,
|
|
},
|
|
}
|
|
task := &model.Midjourney{
|
|
UserId: userID,
|
|
Action: "IMAGINE",
|
|
MjId: "mj-accounting-refund",
|
|
ChannelId: executionChannelID,
|
|
Progress: "0%",
|
|
}
|
|
|
|
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
|
require.NoError(t, err)
|
|
require.True(t, prepared)
|
|
assert.Equal(t, chargedQuota, task.Quota)
|
|
assert.Zero(t, task.TokenId)
|
|
assert.Equal(t, billingChannelID, task.BillingChannelId)
|
|
require.NoError(t, task.Insert())
|
|
|
|
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
|
require.NoError(t, err)
|
|
require.True(t, billed)
|
|
assert.Equal(t, initialUserQuota-chargedQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota-chargedQuota, getTokenRemainQuota(t, tokenID))
|
|
persisted := getMidjourneyTask(t, task.Id)
|
|
assert.Equal(t, chargedQuota, persisted.Quota)
|
|
assert.Equal(t, tokenID, persisted.TokenId)
|
|
assert.Equal(t, billingChannelID, persisted.BillingChannelId)
|
|
|
|
seedChargedAccounting(t, userID, billingChannelID, tokenID, chargedQuota, 1)
|
|
|
|
assert.True(t, RefundMidjourneyQuota(ctx, task, "构图失败"))
|
|
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
|
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, billingChannelID))
|
|
assert.Zero(t, getChannelUsedQuota(t, executionChannelID))
|
|
|
|
persisted = getMidjourneyTask(t, task.Id)
|
|
assert.Zero(t, persisted.Quota)
|
|
assert.Equal(t, tokenID, persisted.TokenId)
|
|
assert.Equal(t, billingChannelID, persisted.BillingChannelId)
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
assert.Equal(t, chargedQuota, log.Quota)
|
|
assert.Equal(t, tokenID, log.TokenId)
|
|
assert.Equal(t, billingChannelID, log.ChannelId)
|
|
|
|
assert.True(t, RefundMidjourneyQuota(ctx, task, "duplicate poll"))
|
|
assert.Equal(t, int64(1), countLogs(t))
|
|
}
|
|
|
|
func TestSettleMidjourneyTaskBillingFundingFailureClearsMarkers(t *testing.T) {
|
|
truncate(t)
|
|
|
|
const userID, tokenID, channelID = 52, 52, 52
|
|
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
|
seedUser(t, userID, initialUserQuota)
|
|
seedToken(t, tokenID, userID, "sk-midjourney-funding-failure", initialTokenQuota)
|
|
seedChannel(t, channelID)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
UserId: userID,
|
|
TokenId: tokenID,
|
|
TokenKey: "sk-midjourney-funding-failure",
|
|
UserQuota: initialUserQuota,
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelId: channelID,
|
|
},
|
|
}
|
|
task := &model.Midjourney{UserId: userID, MjId: "mj-funding-failure", ChannelId: channelID}
|
|
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
|
require.NoError(t, err)
|
|
require.True(t, prepared)
|
|
require.NoError(t, task.Insert())
|
|
|
|
require.NoError(t, model.DB.Exec(`
|
|
CREATE TRIGGER fail_midjourney_user_update
|
|
BEFORE UPDATE ON users
|
|
WHEN OLD.id = 52
|
|
BEGIN
|
|
SELECT RAISE(ABORT, 'forced user quota failure');
|
|
END;
|
|
`).Error)
|
|
t.Cleanup(func() {
|
|
model.DB.Exec("DROP TRIGGER IF EXISTS fail_midjourney_user_update")
|
|
})
|
|
|
|
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
|
|
|
require.Error(t, err)
|
|
assert.False(t, billed)
|
|
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
|
persisted := getMidjourneyTask(t, task.Id)
|
|
assert.Zero(t, persisted.Quota)
|
|
assert.Zero(t, persisted.TokenId)
|
|
assert.Zero(t, persisted.BillingChannelId)
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Zero(t, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
assert.Zero(t, countLogs(t))
|
|
}
|
|
|
|
func TestSettleMidjourneyTaskBillingTokenFailureKeepsFundingRefundable(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 53, 53, 53
|
|
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
|
seedUser(t, userID, initialUserQuota)
|
|
seedToken(t, tokenID, userID, "sk-midjourney-token-failure", initialTokenQuota)
|
|
seedChannel(t, channelID)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
UserId: userID,
|
|
TokenId: tokenID,
|
|
TokenKey: "sk-midjourney-token-failure",
|
|
UserQuota: initialUserQuota,
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelId: channelID,
|
|
},
|
|
}
|
|
task := &model.Midjourney{UserId: userID, MjId: "mj-token-failure", ChannelId: channelID}
|
|
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
|
require.NoError(t, err)
|
|
require.True(t, prepared)
|
|
require.NoError(t, task.Insert())
|
|
|
|
require.NoError(t, model.DB.Exec(`
|
|
CREATE TRIGGER fail_midjourney_token_update
|
|
BEFORE UPDATE ON tokens
|
|
WHEN OLD.id = 53
|
|
BEGIN
|
|
SELECT RAISE(ABORT, 'forced token quota failure');
|
|
END;
|
|
`).Error)
|
|
t.Cleanup(func() {
|
|
model.DB.Exec("DROP TRIGGER IF EXISTS fail_midjourney_token_update")
|
|
})
|
|
|
|
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
|
|
|
require.Error(t, err)
|
|
require.True(t, billed)
|
|
assert.Equal(t, initialUserQuota-chargedQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
|
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
|
persisted := getMidjourneyTask(t, task.Id)
|
|
assert.Equal(t, chargedQuota, persisted.Quota)
|
|
assert.Zero(t, persisted.TokenId)
|
|
assert.Equal(t, channelID, persisted.BillingChannelId)
|
|
|
|
seedChargedAccounting(t, userID, channelID, 0, chargedQuota, 1)
|
|
assert.True(t, RefundMidjourneyQuota(ctx, task, "token settlement failed"))
|
|
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Zero(t, log.TokenId)
|
|
}
|
|
|
|
func TestPrepareMidjourneyTaskBillingRejectsSubscriptionBeforeCharge(t *testing.T) {
|
|
task := &model.Midjourney{Quota: 900, TokenId: 7, BillingChannelId: 8}
|
|
relayInfo := &relaycommon.RelayInfo{BillingSource: BillingSourceSubscription, SubscriptionId: 1}
|
|
|
|
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, 900, true)
|
|
|
|
require.Error(t, err)
|
|
assert.False(t, prepared)
|
|
assert.Zero(t, task.Quota)
|
|
assert.Zero(t, task.TokenId)
|
|
assert.Zero(t, task.BillingChannelId)
|
|
}
|
|
|
|
func TestRefundMidjourneyQuotaUsesLegacyChannelFallbackWithoutTokenAdjustment(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 54, 54, 54
|
|
const walletAfterCharge, tokenQuota, chargedQuota = 7000, 5000, 3000
|
|
seedUser(t, userID, walletAfterCharge)
|
|
seedToken(t, tokenID, userID, "sk-midjourney-legacy", tokenQuota)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, 0, chargedQuota, 1)
|
|
task := &model.Midjourney{
|
|
UserId: userID,
|
|
MjId: "mj-legacy-fallback",
|
|
Action: "IMAGINE",
|
|
ChannelId: channelID,
|
|
Quota: chargedQuota,
|
|
TokenId: 0,
|
|
Progress: "0%",
|
|
}
|
|
require.NoError(t, task.Insert())
|
|
|
|
assert.True(t, RefundMidjourneyQuota(ctx, task, "legacy failure"))
|
|
|
|
assert.Equal(t, walletAfterCharge+chargedQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, tokenQuota, getTokenRemainQuota(t, tokenID))
|
|
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, channelID, log.ChannelId)
|
|
assert.Zero(t, log.TokenId)
|
|
}
|
|
|
|
// ===========================================================================
|
|
// RefundTaskQuota tests
|
|
// ===========================================================================
|
|
|
|
func TestRefundTaskQuota_Wallet(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 1, 1, 1
|
|
const initQuota, preConsumed = 10000, 3000
|
|
const tokenRemain = 5000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-test-key", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
assert.True(t, RefundTaskQuota(ctx, task, "task failed: upstream error"))
|
|
|
|
// User quota should increase by preConsumed
|
|
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
|
|
|
// Token remain_quota should increase, used_quota should decrease
|
|
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
|
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
|
|
// A refund log should be created
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
assert.Equal(t, preConsumed, log.Quota)
|
|
assert.Equal(t, "test-model", log.ModelName)
|
|
assert.Zero(t, task.Quota)
|
|
assert.Zero(t, getTaskQuota(t, task.ID))
|
|
}
|
|
|
|
func TestRefundTaskQuota_Subscription(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID, subID = 2, 2, 2, 1
|
|
const preConsumed = 2000
|
|
const subTotal, subUsed int64 = 100000, 50000
|
|
const tokenRemain = 8000
|
|
|
|
seedUser(t, userID, 0)
|
|
seedToken(t, tokenID, userID, "sk-sub-key", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedSubscription(t, subID, userID, subTotal, subUsed)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
assert.True(t, RefundTaskQuota(ctx, task, "subscription task failed"))
|
|
|
|
// Subscription used should decrease by preConsumed
|
|
assert.Equal(t, subUsed-int64(preConsumed), getSubscriptionUsed(t, subID))
|
|
|
|
// Token should also be refunded
|
|
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
|
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
assert.Zero(t, getTaskQuota(t, task.ID))
|
|
}
|
|
|
|
func TestRefundTaskQuota_ZeroQuota(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID = 3
|
|
seedUser(t, userID, 5000)
|
|
|
|
task := makeTask(userID, 0, 0, 0, BillingSourceWallet, 0)
|
|
|
|
assert.True(t, RefundTaskQuota(ctx, task, "zero quota task"))
|
|
|
|
// No change to user quota
|
|
assert.Equal(t, 5000, getUserQuota(t, userID))
|
|
|
|
// No log created
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestRefundTaskQuota_NoToken(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, channelID = 4, 4
|
|
const initQuota, preConsumed = 10000, 1500
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, 0, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceWallet, 0) // TokenId=0
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
assert.True(t, RefundTaskQuota(ctx, task, "no token task failed"))
|
|
|
|
// User quota refunded
|
|
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
|
|
// Log created
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
assert.Zero(t, getTaskQuota(t, task.ID))
|
|
}
|
|
|
|
func TestRefundTaskQuota_FundingFailureKeepsAccountingAndPendingMarker(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, channelID, preConsumed = 5, 5, 1200
|
|
seedUser(t, userID, 5000)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, 0, preConsumed, 1)
|
|
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceSubscription, 9999)
|
|
task.Status = model.TaskStatusFailure
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
assert.False(t, RefundTaskQuota(ctx, task, "subscription missing"))
|
|
assert.Equal(t, 5000, getUserQuota(t, userID))
|
|
assert.Equal(t, preConsumed, task.Quota)
|
|
assert.Equal(t, preConsumed, getTaskQuota(t, task.ID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, preConsumed, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(preConsumed), getChannelUsedQuota(t, channelID))
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
// ===========================================================================
|
|
// RecalculateTaskQuota tests
|
|
// ===========================================================================
|
|
|
|
func TestRecalculate_PositiveDelta(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 10, 10, 10
|
|
const initQuota, preConsumed = 10000, 2000
|
|
const actualQuota = 3000 // under-charged by 1000
|
|
const tokenRemain = 5000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-recalc-pos", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
|
|
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
|
|
|
|
// User quota should decrease by the delta (1000 additional charge)
|
|
assert.Equal(t, initQuota-(actualQuota-preConsumed), getUserQuota(t, userID))
|
|
|
|
// Token should also be charged the delta
|
|
assert.Equal(t, tokenRemain-(actualQuota-preConsumed), getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, actualQuota, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
|
|
|
// task.Quota should be updated to actualQuota
|
|
assert.Equal(t, actualQuota, task.Quota)
|
|
|
|
// Log type should be Consume (additional charge)
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeConsume, log.Type)
|
|
assert.Equal(t, actualQuota-preConsumed, log.Quota)
|
|
}
|
|
|
|
func TestRecalculate_NegativeDelta(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 11, 11, 11
|
|
const initQuota, preConsumed = 10000, 5000
|
|
const actualQuota = 3000 // over-charged by 2000
|
|
const tokenRemain = 5000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-recalc-neg", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
|
|
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
|
|
|
|
// User quota should increase by abs(delta) = 2000 (refund overpayment)
|
|
assert.Equal(t, initQuota+(preConsumed-actualQuota), getUserQuota(t, userID))
|
|
|
|
// Token should be refunded the difference
|
|
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, actualQuota, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
|
|
|
// task.Quota updated
|
|
assert.Equal(t, actualQuota, task.Quota)
|
|
|
|
// Log type should be Refund
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
assert.Equal(t, preConsumed-actualQuota, log.Quota)
|
|
}
|
|
|
|
func TestRecalculate_ZeroDelta(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID = 12
|
|
const initQuota, preConsumed = 10000, 3000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
|
|
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
|
|
|
RecalculateTaskQuota(ctx, task, preConsumed, "exact match")
|
|
|
|
// No change to user quota
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
|
|
// No log created (delta is zero)
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestRecalculate_ActualQuotaZero(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID = 13
|
|
const initQuota = 10000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
|
|
task := makeTask(userID, 0, 5000, 0, BillingSourceWallet, 0)
|
|
|
|
RecalculateTaskQuota(ctx, task, 0, "zero actual")
|
|
|
|
// No change (early return)
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID, subID = 14, 14, 14, 2
|
|
const preConsumed = 5000
|
|
const actualQuota = 2000 // over-charged by 3000
|
|
const subTotal, subUsed int64 = 100000, 50000
|
|
const tokenRemain = 8000
|
|
|
|
seedUser(t, userID, 0)
|
|
seedToken(t, tokenID, userID, "sk-sub-recalc", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedSubscription(t, subID, userID, subTotal, subUsed)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
|
|
|
|
RecalculateTaskQuota(ctx, task, actualQuota, "subscription over-charge")
|
|
|
|
// Subscription used should decrease by delta (refund 3000)
|
|
assert.Equal(t, subUsed-int64(preConsumed-actualQuota), getSubscriptionUsed(t, subID))
|
|
|
|
// Token refunded
|
|
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, actualQuota, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
|
|
|
assert.Equal(t, actualQuota, task.Quota)
|
|
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
}
|
|
|
|
// ===========================================================================
|
|
// CAS + Billing integration tests
|
|
// Simulates the flow in updateVideoSingleTask (service/task_polling.go)
|
|
// ===========================================================================
|
|
|
|
// simulatePollBilling reproduces the CAS + billing logic from updateVideoSingleTask.
|
|
// It takes a persisted task (already in DB), applies the new status, and performs
|
|
// the conditional update + billing exactly as the polling loop does.
|
|
func simulatePollBilling(ctx context.Context, task *model.Task, newStatus model.TaskStatus, actualQuota int) {
|
|
snap := task.Snapshot()
|
|
|
|
shouldRefund := false
|
|
shouldSettle := false
|
|
quota := task.Quota
|
|
|
|
task.Status = newStatus
|
|
switch string(newStatus) {
|
|
case model.TaskStatusSuccess:
|
|
task.Progress = "100%"
|
|
task.FinishTime = 9999
|
|
shouldSettle = true
|
|
case model.TaskStatusFailure:
|
|
task.Progress = "100%"
|
|
task.FinishTime = 9999
|
|
task.FailReason = "upstream error"
|
|
if quota != 0 {
|
|
shouldRefund = true
|
|
}
|
|
default:
|
|
task.Progress = "50%"
|
|
}
|
|
|
|
isDone := task.Status == model.TaskStatus(model.TaskStatusSuccess) || task.Status == model.TaskStatus(model.TaskStatusFailure)
|
|
if isDone && snap.Status != task.Status {
|
|
won, err := task.UpdateWithStatus(snap.Status)
|
|
if err != nil {
|
|
shouldRefund = false
|
|
shouldSettle = false
|
|
} else if !won {
|
|
shouldRefund = false
|
|
shouldSettle = false
|
|
}
|
|
} else if !snap.Equal(task.Snapshot()) {
|
|
_, _ = task.UpdateWithStatus(snap.Status)
|
|
}
|
|
|
|
if shouldSettle && actualQuota > 0 {
|
|
RecalculateTaskQuota(ctx, task, actualQuota, "test settle")
|
|
}
|
|
if shouldRefund {
|
|
RefundTaskQuota(ctx, task, task.FailReason)
|
|
}
|
|
}
|
|
|
|
func TestCASGuardedRefund_Win(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 20, 20, 20
|
|
const initQuota, preConsumed = 10000, 4000
|
|
const tokenRemain = 6000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-cas-refund-win", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
simulatePollBilling(ctx, task, model.TaskStatus(model.TaskStatusFailure), 0)
|
|
|
|
// CAS wins: task in DB should now be FAILURE
|
|
var reloaded model.Task
|
|
require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
|
|
assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status)
|
|
assert.Zero(t, reloaded.Quota)
|
|
|
|
// Refund should have happened
|
|
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Zero(t, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
|
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
}
|
|
|
|
func TestCASGuardedRefund_Lose(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 21, 21, 21
|
|
const initQuota, preConsumed = 10000, 4000
|
|
const tokenRemain = 6000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-cas-refund-lose", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
// Create task with IN_PROGRESS in DB
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
// Simulate another process already transitioning to FAILURE
|
|
model.DB.Model(&model.Task{}).Where("id = ?", task.ID).Update("status", model.TaskStatusFailure)
|
|
|
|
// Our process still has the old in-memory state (IN_PROGRESS) and tries to transition
|
|
// task.Status is still IN_PROGRESS in the snapshot
|
|
simulatePollBilling(ctx, task, model.TaskStatus(model.TaskStatusFailure), 0)
|
|
|
|
// CAS lost: user quota should NOT change (no double refund)
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, preConsumed, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(preConsumed), getChannelUsedQuota(t, channelID))
|
|
|
|
// No billing log should be created
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestCASGuardedSettle_Win(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 22, 22, 22
|
|
const initQuota, preConsumed = 10000, 5000
|
|
const actualQuota = 3000 // over-charged, should get partial refund
|
|
const tokenRemain = 8000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-cas-settle-win", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
simulatePollBilling(ctx, task, model.TaskStatus(model.TaskStatusSuccess), actualQuota)
|
|
|
|
// CAS wins: task should be SUCCESS
|
|
var reloaded model.Task
|
|
require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
|
|
assert.EqualValues(t, model.TaskStatusSuccess, reloaded.Status)
|
|
|
|
// Settlement should refund the over-charge (5000 - 3000 = 2000 back to user)
|
|
assert.Equal(t, initQuota+(preConsumed-actualQuota), getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
|
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
|
assert.Equal(t, actualQuota, usedQuota)
|
|
assert.Equal(t, 1, requestCount)
|
|
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
|
|
|
// task.Quota should be updated to actualQuota
|
|
assert.Equal(t, actualQuota, task.Quota)
|
|
}
|
|
|
|
func TestNonTerminalUpdate_NoBilling(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, channelID = 23, 23
|
|
const initQuota, preConsumed = 10000, 3000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedChannel(t, channelID)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceWallet, 0)
|
|
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
|
task.Progress = "20%"
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
|
|
// Simulate a non-terminal poll update (still IN_PROGRESS, progress changed)
|
|
simulatePollBilling(ctx, task, model.TaskStatus(model.TaskStatusInProgress), 0)
|
|
|
|
// User quota should NOT change
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
|
|
// No billing log
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
|
|
// Task progress should be updated in DB
|
|
var reloaded model.Task
|
|
require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
|
|
assert.Equal(t, "50%", reloaded.Progress)
|
|
}
|
|
|
|
// ===========================================================================
|
|
// Mock adaptor for settleTaskBillingOnComplete tests
|
|
// ===========================================================================
|
|
|
|
type mockAdaptor struct {
|
|
adjustReturn int
|
|
}
|
|
|
|
func (m *mockAdaptor) Init(_ *relaycommon.RelayInfo) {}
|
|
func (m *mockAdaptor) FetchTask(string, string, map[string]any, string) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
func (m *mockAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { return nil, nil }
|
|
func (m *mockAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
|
|
return m.adjustReturn
|
|
}
|
|
|
|
// ===========================================================================
|
|
// PerCallBilling tests — settleTaskBillingOnComplete
|
|
// ===========================================================================
|
|
|
|
func TestSettle_PerCallBilling_SkipsAdaptorAdjust(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 30, 30, 30
|
|
const initQuota, preConsumed = 10000, 5000
|
|
const tokenRemain = 8000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-percall-adaptor", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
task.PrivateData.BillingContext.PerCallBilling = true
|
|
|
|
adaptor := &mockAdaptor{adjustReturn: 2000}
|
|
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
|
|
|
|
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
|
|
|
// Per-call: no adjustment despite adaptor returning 2000
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, preConsumed, task.Quota)
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestSettle_PerCallBilling_SkipsTotalTokens(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 31, 31, 31
|
|
const initQuota, preConsumed = 10000, 4000
|
|
const tokenRemain = 7000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-percall-tokens", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
task.PrivateData.BillingContext.PerCallBilling = true
|
|
|
|
adaptor := &mockAdaptor{adjustReturn: 0}
|
|
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, TotalTokens: 9999}
|
|
|
|
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
|
|
|
// Per-call: no recalculation by tokens
|
|
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, preConsumed, task.Quota)
|
|
assert.Equal(t, int64(0), countLogs(t))
|
|
}
|
|
|
|
func TestSettle_NonPerCallBilling_AppliesAdaptorAdjustment(t *testing.T) {
|
|
truncate(t)
|
|
ctx := context.Background()
|
|
|
|
const userID, tokenID, channelID = 32, 32, 32
|
|
const initQuota, preConsumed = 10000, 5000
|
|
const adaptorQuota = 3000
|
|
const tokenRemain = 8000
|
|
|
|
seedUser(t, userID, initQuota)
|
|
seedToken(t, tokenID, userID, "sk-nonpercall-adj", tokenRemain)
|
|
seedChannel(t, channelID)
|
|
|
|
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
|
// PerCallBilling defaults to false
|
|
|
|
adaptor := &mockAdaptor{adjustReturn: adaptorQuota}
|
|
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
|
|
|
|
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
|
|
|
// Non-per-call: adaptor adjustment applies (refund 2000)
|
|
assert.Equal(t, initQuota+(preConsumed-adaptorQuota), getUserQuota(t, userID))
|
|
assert.Equal(t, tokenRemain+(preConsumed-adaptorQuota), getTokenRemainQuota(t, tokenID))
|
|
assert.Equal(t, adaptorQuota, task.Quota)
|
|
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Equal(t, model.LogTypeRefund, log.Type)
|
|
}
|