cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

feat: publish sanitized custom.97 source

Browse Source
This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

service/task_billing_test.go
+77 -11
View File
@@ -13,8 +13,10 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
@@ -44,6 +46,9 @@ func TestMain(m *testing.M) {
common.RedisEnabled = false
common.BatchUpdateEnabled = false
common.LogConsumeEnabled = true
if err := i18n.Init(); err != nil {
panic("failed to load locales: " + err.Error())
}
if err := db.AutoMigrate(
&model.Task{},
@@ -341,7 +346,7 @@ func TestLogTaskConsumptionIncludesTieredSnapshotUsageFacts(t *testing.T) {
assert.Equal(t, float64(5), facts["seconds"])
assert.NotContains(t, other, "resolution")
assert.NotContains(t, other, "seconds")
assert.Contains(t, log.Content, "计算参数：")
assert.Contains(t, log.Content, "Billing parameters: ")
assert.Contains(t, log.Content, "resolution: 720P")
assert.Contains(t, log.Content, "seconds: 5")
}
@@ -379,8 +384,69 @@ func TestLogTaskConsumptionWithoutSnapshotKeepsRatioMode(t *testing.T) {
assert.NotContains(t, other, "expr_b64")
assert.NotContains(t, other, "matched_tier")
assert.NotContains(t, other, "usage_facts")
assert.Contains(t, log.Content, "计算参数：")
assert.Contains(t, log.Content, "size: 2.00")
assert.Equal(t, "Action GENERATE, Billing parameters: size: 2.00", log.Content)
// The web console renders content_parts in the viewer's language; content
// keeps the English text.
assert.Equal(t, []any{
map[string]any{"key": "Action {{action}}", "params": map[string]any{"action": "GENERATE"}},
map[string]any{"key": "Billing parameters: {{params}}", "params": map[string]any{"params": "size: 2.00"}},
}, other["content_parts"])
}
// Task logs distinguish jobs the client polls from requests whose HTTP call
// returned the deliverable itself, and flag results the gateway did not keep.
func TestLogTaskConsumptionMarksInlineResultsAndDiscardedArtifacts(t *testing.T) {
for _, tc := range []struct {
name string
status model.TaskStatus
discarded bool
pinnedProtocol string
wantSync, wantKept bool
}{
{"asynchronous job", model.TaskStatusNotStart, false, "", false, true},
{"immediate result on a discarding route", model.TaskStatusSuccess, true, "", true, false},
{"openai image request waits for an asynchronous upstream task", model.TaskStatusNotStart, false, jsplugin.ProtocolOpenAIImage, true, true},
{"immediate failure", model.TaskStatusFailure, false, "", true, true},
} {
t.Run(tc.name, func(t *testing.T) {
truncate(t)
const userID, channelID = 43, 43
seedUser(t, userID, 10_000)
seedChannel(t, channelID)
task := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0)
task.Status = tc.status
task.PrivateData.ResultDiscarded = tc.discarded
info := &relaycommon.RelayInfo{
UserId: userID, OriginModelName: "qwen-image-plus", UsingGroup: "default",
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "text_to_image"},
PriceData: types.PriceData{ModelPrice: 0.03, Quota: 100, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}},
}
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
ctx.Set("token_name", "test_token")
if tc.pinnedProtocol != "" {
ctx.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Protocol: tc.pinnedProtocol})
}
LogTaskConsumption(ctx, info, task)
log := getLastLog(t)
require.NotNil(t, log)
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
assert.Equal(t, true, other["is_task"])
if tc.wantSync {
assert.Equal(t, true, other["task_sync"])
} else {
assert.NotContains(t, other, "task_sync")
}
if tc.wantKept {
assert.NotContains(t, other, "result_discarded")
} else {
assert.Equal(t, true, other["result_discarded"])
}
})
}
}
func TestTaskBillingOtherSeparatesPluginAndRootDiagnostics(t *testing.T) {
@@ -621,7 +687,7 @@ func TestMidjourneyRefundRestoresEveryAccountingElementOnBillingChannel(t *testi
seedChargedAccounting(t, userID, billingChannelID, tokenID, chargedQuota, 1)
assert.True(t, RefundMidjourneyQuota(ctx, task, "构图失败"))
assert.True(t, RefundMidjourneyQuota(ctx, task, "Composition failed"))
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
assert.Zero(t, getTokenUsedQuota(t, tokenID))
@@ -981,7 +1047,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("adaptor adjustment"))
// User quota should decrease by the delta (1000 additional charge)
assert.Equal(t, initQuota-(actualQuota-preConsumed), getUserQuota(t, userID))
@@ -1020,7 +1086,7 @@ func TestRecalculate_NegativeDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("adaptor adjustment"))
// User quota should increase by abs(delta) = 2000 (refund overpayment)
assert.Equal(t, initQuota+(preConsumed-actualQuota), getUserQuota(t, userID))
@@ -1054,7 +1120,7 @@ func TestRecalculate_ZeroDelta(t *testing.T) {
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, preConsumed, "exact match")
RecalculateTaskQuota(ctx, task, preConsumed, common.NewMessage("exact match"))
// No change to user quota
assert.Equal(t, initQuota, getUserQuota(t, userID))
@@ -1075,7 +1141,7 @@ func TestRecalculate_ActualQuotaZero(t *testing.T) {
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
require.NoError(t, model.DB.Create(task).Error)
RecalculateTaskQuota(ctx, task, 0, "zero actual")
RecalculateTaskQuota(ctx, task, 0, common.NewMessage("zero actual"))
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
assert.Zero(t, task.Quota)
@@ -1094,7 +1160,7 @@ func TestRecalculate_RejectsNegativeActualQuota(t *testing.T) {
seedUser(t, userID, initQuota)
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, -1, "invalid negative actual")
RecalculateTaskQuota(ctx, task, -1, common.NewMessage("invalid negative actual"))
assert.Equal(t, initQuota, getUserQuota(t, userID))
assert.Equal(t, preConsumed, task.Quota)
@@ -1119,7 +1185,7 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
RecalculateTaskQuota(ctx, task, actualQuota, "subscription over-charge")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("subscription over-charge"))
// Subscription used should decrease by delta (refund 3000)
assert.Equal(t, subUsed-int64(preConsumed-actualQuota), getSubscriptionUsed(t, subID))
@@ -1186,7 +1252,7 @@ func simulatePollBilling(ctx context.Context, task *model.Task, newStatus model.
}
if shouldSettle && actualQuota > 0 {
RecalculateTaskQuota(ctx, task, actualQuota, "test settle")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("test settle"))
}
if shouldRefund {
RefundTaskQuota(ctx, task, task.FailReason)