diff --git a/controller/midjourney.go b/controller/midjourney.go index bf52314a..0494386c 100644 --- a/controller/midjourney.go +++ b/controller/midjourney.go @@ -213,22 +213,7 @@ func runMidjourneyTaskUpdateOnce(ctx context.Context, report func(processed, tot if err != nil { logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) } else if won && shouldReturnQuota { - err = model.IncreaseUserQuota(task.UserId, task.Quota, false) - if err != nil { - logger.LogError(ctx, "fail to increase user quota: "+err.Error()) - } - model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ - UserId: task.UserId, - LogType: model.LogTypeRefund, - Content: "", - ChannelId: task.ChannelId, - ModelName: service.CovertMjpActionToModelName(task.Action), - Quota: task.Quota, - Other: map[string]interface{}{ - "task_id": task.MjId, - "reason": "构图失败", - }, - }) + service.RefundMidjourneyQuota(ctx, task, "构图失败") } } } diff --git a/model/midjourney.go b/model/midjourney.go index 201f774c..e42c25e0 100644 --- a/model/midjourney.go +++ b/model/midjourney.go @@ -23,6 +23,9 @@ type Midjourney struct { Quota int `json:"quota"` Buttons string `json:"buttons"` Properties string `json:"properties"` + + TokenId int `json:"-" gorm:"default:0"` + BillingChannelId int `json:"-" gorm:"default:0"` } // TaskQueryParams 用于包含所有搜索条件的结构体，可以根据需求添加更多字段 @@ -170,6 +173,19 @@ func (midjourney *Midjourney) Update() error { return err } +func (midjourney *Midjourney) UpdateBillingState() error { + return DB.Model(midjourney). + Select("quota", "token_id", "billing_channel_id"). + Updates(midjourney).Error +} + +func (midjourney *Midjourney) GetBillingChannelId() int { + if midjourney.BillingChannelId > 0 { + return midjourney.BillingChannelId + } + return midjourney.ChannelId +} + // UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS). // Returns (true, nil) if this caller won the update, (false, nil) if // another process already moved the task out of fromStatus. diff --git a/model/user.go b/model/user.go index 0a9a42a7..7bc060ad 100644 --- a/model/user.go +++ b/model/user.go @@ -1353,6 +1353,17 @@ func UpdateUserUsedQuotaAndRequestCount(id int, quota int) { updateUserUsedQuotaAndRequestCount(id, quota, 1) } +// UpdateUserUsedQuota adjusts accumulated usage without changing request count. +func UpdateUserUsedQuota(id int, quota int) { + if common.BatchUpdateEnabled { + addNewRecord(BatchUpdateTypeUsedQuota, id, quota) + return + } + if err := DB.Model(&User{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error; err != nil { + common.SysLog("failed to update user used quota: " + err.Error()) + } +} + func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) { err := DB.Model(&User{}).Where("id = ?", id).Updates( map[string]interface{}{ diff --git a/model/user_update_test.go b/model/user_update_test.go index 1a89e93a..a0b6d914 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -89,6 +89,61 @@ func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing assert.Equal(t, "rotated-token", got.GetAccessToken()) } +func TestUsageAccountingSupportsSignedDirectAndBatchDeltas(t *testing.T) { + setupUserUpdateTestState(t) + resetBatchUpdateTestState(t) + + user := User{ + Id: 10, + Username: "usage-adjustment-user", + Password: "password", + Status: common.UserStatusEnabled, + UsedQuota: 1000, + RequestCount: 3, + } + channel := Channel{ + Id: 10, + Name: "usage-adjustment-channel", + Key: "sk-test", + Status: common.ChannelStatusEnabled, + UsedQuota: 1000, + } + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Create(&channel).Error) + + UpdateUserUsedQuota(user.Id, -200) + UpdateUserUsedQuota(user.Id, 50) + UpdateChannelUsedQuota(channel.Id, -200) + UpdateChannelUsedQuota(channel.Id, 50) + + var got User + require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error) + assert.Equal(t, 850, got.UsedQuota) + assert.Equal(t, 3, got.RequestCount) + var gotChannel Channel + require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error) + assert.Equal(t, int64(850), gotChannel.UsedQuota) + + common.BatchUpdateEnabled = true + UpdateUserUsedQuota(user.Id, 400) + UpdateUserUsedQuota(user.Id, -100) + UpdateChannelUsedQuota(channel.Id, 400) + UpdateChannelUsedQuota(channel.Id, -100) + + require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error) + assert.Equal(t, 850, got.UsedQuota, "batch deltas must remain queued until flush") + assert.Equal(t, 3, got.RequestCount) + require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error) + assert.Equal(t, int64(850), gotChannel.UsedQuota, "batch deltas must remain queued until flush") + + batchUpdate() + require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error) + assert.Equal(t, 1150, got.UsedQuota) + assert.Equal(t, 3, got.RequestCount) + require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error) + assert.Equal(t, int64(1150), gotChannel.UsedQuota) +} + func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) { setupUserUpdateTestState(t) diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index daa0585a..99755f0f 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -232,30 +232,6 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR if err != nil { return &mjResp.Response } - defer func() { - if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 { - err := service.PostConsumeQuota(info, priceData.Quota, 0, true) - if err != nil { - common.SysLog("error consuming token remain quota: " + err.Error()) - } - - tokenName := c.GetString("token_name") - logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace) - other := service.GenerateMjOtherInfo(info, priceData) - model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ - ChannelId: info.ChannelId, - ModelName: modelName, - TokenName: tokenName, - Quota: priceData.Quota, - Content: logContent, - TokenId: info.TokenId, - Group: info.UsingGroup, - Other: other, - }) - model.UpdateUserUsedQuotaAndRequestCount(info.UserId, priceData.Quota) - model.UpdateChannelUsedQuota(info.ChannelId, priceData.Quota) - } - }() midjResponse := &mjResp.Response midjourneyTask := &model.Midjourney{ UserId: info.UserId, @@ -274,12 +250,42 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR Progress: "0%", FailReason: "", ChannelId: c.GetInt("channel_id"), - Quota: priceData.Quota, + } + billingPrepared, billingErr := service.PrepareMidjourneyTaskBilling( + info, + midjourneyTask, + priceData.Quota, + mjResp.StatusCode == http.StatusOK && midjResponse.Code == 1, + ) + if billingErr != nil { + common.SysLog("error consuming Midjourney quota: " + billingErr.Error()) } err = midjourneyTask.Insert() if err != nil { return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed") } + billingApplied, billingErr := service.SettleMidjourneyTaskBilling(info, midjourneyTask, billingPrepared) + if billingErr != nil { + common.SysLog("error settling Midjourney quota: " + billingErr.Error()) + } + if billingApplied { + billingChannelId := midjourneyTask.GetBillingChannelId() + tokenName := c.GetString("token_name") + logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace) + other := service.GenerateMjOtherInfo(info, priceData) + model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ + ChannelId: billingChannelId, + ModelName: modelName, + TokenName: tokenName, + Quota: midjourneyTask.Quota, + Content: logContent, + TokenId: midjourneyTask.TokenId, + Group: info.UsingGroup, + Other: other, + }) + model.UpdateUserUsedQuotaAndRequestCount(info.UserId, midjourneyTask.Quota) + model.UpdateChannelUsedQuota(billingChannelId, midjourneyTask.Quota) + } c.Writer.WriteHeader(mjResp.StatusCode) respBody, err := json.Marshal(midjResponse) if err != nil { @@ -539,30 +545,6 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt } midjResponse := &midjResponseWithStatus.Response - defer func() { - if consumeQuota && midjResponseWithStatus.StatusCode == 200 { - err := service.PostConsumeQuota(relayInfo, priceData.Quota, 0, true) - if err != nil { - common.SysLog("error consuming token remain quota: " + err.Error()) - } - tokenName := c.GetString("token_name") - logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s，ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result) - other := service.GenerateMjOtherInfo(relayInfo, priceData) - model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{ - ChannelId: relayInfo.ChannelId, - ModelName: modelName, - TokenName: tokenName, - Quota: priceData.Quota, - Content: logContent, - TokenId: relayInfo.TokenId, - Group: relayInfo.UsingGroup, - Other: other, - }) - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, priceData.Quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, priceData.Quota) - } - }() - // 文档：https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md //1-提交成功 // 21-任务已存在（处理中或者有结果了） {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}} @@ -587,7 +569,6 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt Progress: "0%", FailReason: "", ChannelId: c.GetInt("channel_id"), - Quota: priceData.Quota, } if midjResponse.Code == 3 { //无实例账号自动禁用渠道（No available account instance） @@ -632,6 +613,15 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt midjourneyTask.Progress = "100%" midjourneyTask.Status = "SUCCESS" } + billingPrepared, billingErr := service.PrepareMidjourneyTaskBilling( + relayInfo, + midjourneyTask, + priceData.Quota, + consumeQuota && midjResponseWithStatus.StatusCode == http.StatusOK, + ) + if billingErr != nil { + common.SysLog("error consuming Midjourney quota: " + billingErr.Error()) + } err = midjourneyTask.Insert() if err != nil { return &dto.MidjourneyResponse{ @@ -639,6 +629,28 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt Description: "insert_midjourney_task_failed", } } + billingApplied, billingErr := service.SettleMidjourneyTaskBilling(relayInfo, midjourneyTask, billingPrepared) + if billingErr != nil { + common.SysLog("error settling Midjourney quota: " + billingErr.Error()) + } + if billingApplied { + billingChannelId := midjourneyTask.GetBillingChannelId() + tokenName := c.GetString("token_name") + logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s，ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result) + other := service.GenerateMjOtherInfo(relayInfo, priceData) + model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{ + ChannelId: billingChannelId, + ModelName: modelName, + TokenName: tokenName, + Quota: midjourneyTask.Quota, + Content: logContent, + TokenId: midjourneyTask.TokenId, + Group: relayInfo.UsingGroup, + Other: other, + }) + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, midjourneyTask.Quota) + model.UpdateChannelUsedQuota(billingChannelId, midjourneyTask.Quota) + } if midjResponse.Code == 22 { //22-排队中，说明任务已存在 //修改返回值 diff --git a/service/midjourney.go b/service/midjourney.go index c6397b7f..b10f6cd2 100644 --- a/service/midjourney.go +++ b/service/midjourney.go @@ -3,6 +3,8 @@ package service import ( "context" "encoding/json" + "errors" + "fmt" "io" "net/http" "strconv" @@ -13,6 +15,8 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/setting" @@ -27,6 +31,113 @@ func CovertMjpActionToModelName(mjAction string) string { return modelName } +// PrepareMidjourneyTaskBilling sets the durable refund marker before the task is inserted. +func PrepareMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, quota int, shouldBill bool) (bool, error) { + if task == nil { + return false, errors.New("Midjourney task is nil") + } + task.Quota = 0 + task.TokenId = 0 + task.BillingChannelId = 0 + if !shouldBill { + return false, nil + } + if relayInfo == nil { + return false, errors.New("relay info is nil") + } + if quota < 0 { + return false, errors.New("quota cannot be negative") + } + if relayInfo.BillingSource == BillingSourceSubscription { + return false, errors.New("legacy Midjourney billing does not support subscriptions") + } + + task.Quota = quota + task.BillingChannelId = task.ChannelId + if relayInfo.ChannelMeta != nil && relayInfo.ChannelId > 0 { + task.BillingChannelId = relayInfo.ChannelId + } + return true, nil +} + +// SettleMidjourneyTaskBilling charges a persisted legacy task and records the applied stages. +func SettleMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, prepared bool) (bool, error) { + if !prepared { + return false, nil + } + if relayInfo == nil { + return false, errors.New("relay info is nil") + } + if task == nil || task.Id == 0 { + return false, errors.New("Midjourney task must be persisted before billing") + } + + result, billingErr := postConsumeQuotaWithResult(relayInfo, task.Quota, 0, true) + if !result.FundingApplied { + task.Quota = 0 + task.TokenId = 0 + task.BillingChannelId = 0 + if updateErr := task.UpdateBillingState(); updateErr != nil { + return false, errors.Join(billingErr, fmt.Errorf("clear Midjourney billing state: %w", updateErr)) + } + return false, billingErr + } + + task.TokenId = 0 + if result.TokenApplied { + task.TokenId = relayInfo.TokenId + } + if updateErr := task.UpdateBillingState(); updateErr != nil { + return true, errors.Join(billingErr, fmt.Errorf("update Midjourney billing state: %w", updateErr)) + } + return true, billingErr +} + +// RefundMidjourneyQuota reverses every accounting element recorded for a billed legacy task. +func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason string) bool { + quota := task.Quota + if quota == 0 { + return true + } + + if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil { + logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 用户额度失败 task %s: %s", task.MjId, err.Error())) + return false + } + + if task.TokenId > 0 { + tokenKey := resolveTokenKey(ctx, task.TokenId, task.MjId) + if tokenKey != "" { + if err := model.IncreaseTokenQuota(task.TokenId, tokenKey, quota); err != nil { + logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 令牌额度失败 task %s: %s", task.MjId, err.Error())) + } + } + } + + billingChannelId := task.GetBillingChannelId() + model.UpdateUserUsedQuota(task.UserId, -quota) + model.UpdateChannelUsedQuota(billingChannelId, -quota) + model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ + UserId: task.UserId, + LogType: model.LogTypeRefund, + Content: "", + ChannelId: billingChannelId, + ModelName: CovertMjpActionToModelName(task.Action), + Quota: quota, + TokenId: task.TokenId, + Other: map[string]interface{}{ + "task_id": task.MjId, + "reason": reason, + }, + }) + + task.Quota = 0 + if err := task.UpdateBillingState(); err != nil { + logger.LogError(ctx, fmt.Sprintf("Midjourney 退款成功但清除 quota 失败 task %s: %s", task.MjId, err.Error())) + } + return true +} + func GetMjRequestModel(relayMode int, midjRequest *dto.MidjourneyRequest) (string, *dto.MidjourneyResponse, bool) { action := "" if relayMode == relayconstant.RelayModeMidjourneyAction { diff --git a/service/quota.go b/service/quota.go index 0c0c2da4..359956a5 100644 --- a/service/quota.go +++ b/service/quota.go @@ -406,17 +406,27 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error { return nil } -func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) { +type postConsumeQuotaResult struct { + FundingApplied bool + TokenApplied bool +} + +func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) error { + _, err := postConsumeQuotaWithResult(relayInfo, quota, preConsumedQuota, sendEmail) + return err +} + +func postConsumeQuotaWithResult(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (result postConsumeQuotaResult, err error) { // 1) Consume from wallet quota OR subscription item if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription { if relayInfo.SubscriptionId == 0 { - return errors.New("subscription id is missing") + return result, errors.New("subscription id is missing") } delta := int64(quota) if delta != 0 { if err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta); err != nil { - return err + return result, err } relayInfo.SubscriptionPostDelta += delta } @@ -428,9 +438,10 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false) } if err != nil { - return err + return result, err } } + result.FundingApplied = true if !relayInfo.IsPlayground { if quota > 0 { @@ -439,8 +450,9 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota) } if err != nil { - return err + return result, err } + result.TokenApplied = true } if sendEmail { @@ -449,7 +461,7 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu } } - return nil + return result, nil } func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) { diff --git a/service/task_billing.go b/service/task_billing.go index b677cb62..64dbdcef 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -161,7 +161,7 @@ func taskModelName(task *model.Task) string { } // RefundTaskQuota 统一的任务失败退款逻辑。 -// 当异步任务失败时，将预扣的 quota 退还给用户（支持钱包和订阅），并退还令牌额度。 +// 当异步任务失败时，退还资金与令牌额度，并回减用户和渠道用量。 // 返回资金来源是否已成功退还；失败时保留 quota，供显式重试或人工对账。 func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool { quota := task.Quota @@ -178,7 +178,11 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool // 2. 退还令牌额度 taskAdjustTokenQuota(ctx, task, -quota) - // 3. 记录日志 + // 3. 回减预扣时累计的用户和渠道用量，请求次数保持不变 + model.UpdateUserUsedQuota(task.UserId, -quota) + model.UpdateChannelUsedQuota(task.ChannelId, -quota) + + // 4. 记录日志 other := taskBillingOther(task) other["task_id"] = task.TaskID other["reason"] = reason @@ -194,7 +198,7 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool Other: other, }) - // 4. 资金退款完成后再清除持久化标记。 + // 5. 资金退款完成后再清除持久化标记。 // 回写失败必须显式告警，避免漏掉潜在的重复退款风险。 task.Quota = 0 if err := task.UpdateQuota(); err != nil { @@ -242,13 +246,15 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int logger.LogError(ctx, fmt.Sprintf("差额结算回写 quota 失败 task %s: %s", task.TaskID, err.Error())) } + // 提交阶段已经累计过一次请求；结算阶段只调整最终用量。 + model.UpdateUserUsedQuota(task.UserId, quotaDelta) + model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta) + var logType int var logQuota int if quotaDelta > 0 { logType = model.LogTypeConsume logQuota = quotaDelta - model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta) - model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta) } else { logType = model.LogTypeRefund logQuota = -quotaDelta diff --git a/service/task_billing_test.go b/service/task_billing_test.go index 53e3f680..699cc1ed 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -45,6 +45,7 @@ func TestMain(m *testing.M) { &model.Token{}, &model.Log{}, &model.Channel{}, + &model.Midjourney{}, &model.TopUp{}, &model.UserSubscription{}, &model.SystemTask{}, @@ -68,6 +69,7 @@ func truncate(t *testing.T) { 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") @@ -115,6 +117,20 @@ func seedChannel(t *testing.T, id int) { 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"), @@ -249,6 +265,20 @@ func getUserQuota(t *testing.T, id int) int { 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 @@ -277,6 +307,13 @@ func getTaskQuota(t *testing.T, id int64) int { 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 @@ -294,6 +331,291 @@ func countLogs(t *testing.T) int64 { 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 // =========================================================================== @@ -309,6 +631,7 @@ func TestRefundTaskQuota_Wallet(t *testing.T) { 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) @@ -320,7 +643,11 @@ func TestRefundTaskQuota_Wallet(t *testing.T) { // Token remain_quota should increase, used_quota should decrease assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID)) - assert.Equal(t, -preConsumed, getTokenUsedQuota(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) @@ -345,6 +672,7 @@ func TestRefundTaskQuota_Subscription(t *testing.T) { 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) @@ -356,6 +684,11 @@ func TestRefundTaskQuota_Subscription(t *testing.T) { // 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) @@ -390,6 +723,7 @@ func TestRefundTaskQuota_NoToken(t *testing.T) { 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) @@ -398,6 +732,10 @@ func TestRefundTaskQuota_NoToken(t *testing.T) { // 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) @@ -406,19 +744,26 @@ func TestRefundTaskQuota_NoToken(t *testing.T) { assert.Zero(t, getTaskQuota(t, task.ID)) } -func TestRefundTaskQuota_FundingFailureKeepsPendingMarker(t *testing.T) { +func TestRefundTaskQuota_FundingFailureKeepsAccountingAndPendingMarker(t *testing.T) { truncate(t) ctx := context.Background() - const userID, preConsumed = 5, 1200 + const userID, channelID, preConsumed = 5, 5, 1200 seedUser(t, userID, 5000) - task := makeTask(userID, 0, preConsumed, 0, BillingSourceSubscription, 9999) + 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)) } @@ -438,6 +783,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) { 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) @@ -448,6 +794,11 @@ func TestRecalculate_PositiveDelta(t *testing.T) { // 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) @@ -471,6 +822,7 @@ func TestRecalculate_NegativeDelta(t *testing.T) { 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) @@ -481,6 +833,11 @@ func TestRecalculate_NegativeDelta(t *testing.T) { // 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) @@ -544,6 +901,7 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) { 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) @@ -554,6 +912,11 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) { // 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) @@ -627,6 +990,7 @@ func TestCASGuardedRefund_Win(t *testing.T) { 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) @@ -643,6 +1007,10 @@ func TestCASGuardedRefund_Win(t *testing.T) { // 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) @@ -660,6 +1028,7 @@ func TestCASGuardedRefund_Lose(t *testing.T) { 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) @@ -676,6 +1045,10 @@ func TestCASGuardedRefund_Lose(t *testing.T) { // 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)) @@ -693,6 +1066,7 @@ func TestCASGuardedSettle_Win(t *testing.T) { 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) @@ -708,6 +1082,10 @@ func TestCASGuardedSettle_Win(t *testing.T) { // 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)