feat: publish v1.0.0-rc.36-custom.1
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 9m20s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 55m7s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 9m20s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 55m7s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
This commit is contained in:
Codex committed
commit
bb7968b612
2381 files changed
+551494
No files matched your search
@@ -0,0 +1,999 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
taskdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type taskPollingFetchAdaptor struct {
|
||||
mu sync.Mutex
|
||||
taskIDs []string
|
||||
fetched chan string
|
||||
blockTaskID string
|
||||
blockStarted chan struct{}
|
||||
releaseBlock chan struct{}
|
||||
blockOnce sync.Once
|
||||
}
|
||||
|
||||
type batchPollingAdaptor struct {
|
||||
taskPollingFetchAdaptor
|
||||
batchCalls int
|
||||
batchIDs []string
|
||||
results map[string]*BatchTaskResult
|
||||
}
|
||||
|
||||
func (a *batchPollingAdaptor) FetchMode() string { return "batch" }
|
||||
func (a *batchPollingAdaptor) FetchBatchTasks(_ string, _ string, tasks []*model.Task, _ string) (*http.Response, error) {
|
||||
a.batchCalls++
|
||||
a.batchIDs = a.batchIDs[:0]
|
||||
for _, task := range tasks {
|
||||
a.batchIDs = append(a.batchIDs, task.GetUpstreamTaskID())
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader([]byte(`{}`)))}, nil
|
||||
}
|
||||
func (a *batchPollingAdaptor) ParseBatchResult(_ []*model.Task, _ *http.Response, _ []byte) (map[string]*BatchTaskResult, error) {
|
||||
if a.results != nil {
|
||||
return a.results, nil
|
||||
}
|
||||
results := make(map[string]*BatchTaskResult, len(a.batchIDs))
|
||||
for _, taskID := range a.batchIDs {
|
||||
results[taskID] = &BatchTaskResult{TaskInfo: relaycommon.TaskInfo{TaskID: taskID, Status: model.TaskStatusInProgress, Url: "https://example.com/result"}}
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, task *model.Task, _ string) (*http.Response, error) {
|
||||
taskID := ""
|
||||
if task != nil {
|
||||
taskID = task.GetUpstreamTaskID()
|
||||
}
|
||||
if taskID == a.blockTaskID && a.releaseBlock != nil {
|
||||
a.blockOnce.Do(func() {
|
||||
if a.blockStarted != nil {
|
||||
close(a.blockStarted)
|
||||
}
|
||||
})
|
||||
<-a.releaseBlock
|
||||
}
|
||||
|
||||
a.mu.Lock()
|
||||
a.taskIDs = append(a.taskIDs, taskID)
|
||||
a.mu.Unlock()
|
||||
if a.fetched != nil {
|
||||
select {
|
||||
case a.fetched <- taskID:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
response := taskdto.TaskResponse[model.Task]{
|
||||
Code: taskdto.TaskSuccessCode,
|
||||
Data: model.Task{
|
||||
TaskID: taskID,
|
||||
Status: model.TaskStatusInProgress,
|
||||
Progress: "30%",
|
||||
},
|
||||
}
|
||||
responseBody, err := common.Marshal(response)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(responseBody)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
|
||||
return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) fetchCount() int {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return len(a.taskIDs)
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) fetchedTaskIDs() []string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return append([]string(nil), a.taskIDs...)
|
||||
}
|
||||
|
||||
func TestRedactVideoResponseBodyPreservesPollingPayloadShape(t *testing.T) {
|
||||
rawVideo := strings.Repeat("a", 300)
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"done": true,
|
||||
"name": "operations/provider-task",
|
||||
"response": map[string]any{
|
||||
"bytesBase64Encoded": "secret-bytes",
|
||||
"video": rawVideo,
|
||||
"videos": []any{
|
||||
map[string]any{
|
||||
"bytesBase64Encoded": "other-secret-bytes",
|
||||
"mimeType": "video/mp4",
|
||||
"uri": "https://media.example/video.mp4",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var stored map[string]any
|
||||
require.NoError(t, common.Unmarshal(redactVideoResponseBody(body), &stored))
|
||||
assert.Equal(t, true, stored["done"])
|
||||
assert.Equal(t, "operations/provider-task", stored["name"])
|
||||
response, ok := stored["response"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.NotContains(t, response, "bytesBase64Encoded")
|
||||
assert.Equal(t, strings.Repeat("a", 256)+"...", response["video"])
|
||||
videos, ok := response["videos"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, videos, 1)
|
||||
video, ok := videos[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.NotContains(t, video, "bytesBase64Encoded")
|
||||
assert.Equal(t, "video/mp4", video["mimeType"])
|
||||
assert.Equal(t, "https://media.example/video.mp4", video["uri"])
|
||||
}
|
||||
|
||||
func seedTaskPollingChannel(t *testing.T, id int, disableSleep bool) {
|
||||
t.Helper()
|
||||
ch := &model.Channel{
|
||||
Id: id,
|
||||
Type: constant.ChannelTypeKling,
|
||||
Name: "polling_channel",
|
||||
Key: "sk-test",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
}
|
||||
if disableSleep {
|
||||
ch.SetOtherSettings(dto.ChannelOtherSettings{DisableTaskPollingSleep: true})
|
||||
}
|
||||
require.NoError(t, model.DB.Create(ch).Error)
|
||||
}
|
||||
|
||||
func seedPollingTask(t *testing.T, channelID int, publicID string, upstreamID string) *model.Task {
|
||||
t.Helper()
|
||||
task := &model.Task{
|
||||
TaskID: publicID,
|
||||
Platform: constant.TaskPlatform("kling"),
|
||||
UserId: 1,
|
||||
ChannelId: channelID,
|
||||
Action: constant.TaskActionImageToVideo,
|
||||
Status: model.TaskStatusInProgress,
|
||||
Progress: "30%",
|
||||
CreatedAt: time.Now().Unix(),
|
||||
UpdatedAt: time.Now().Unix(),
|
||||
PrivateData: model.TaskPrivateData{
|
||||
UpstreamTaskID: upstreamID,
|
||||
},
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
return task
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksDefaultSleepWaitsBetweenTasks(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const channelID = 101
|
||||
seedTaskPollingChannel(t, channelID, false)
|
||||
first := seedPollingTask(t, channelID, "task_public_1", "upstream_1")
|
||||
second := seedPollingTask(t, channelID, "task_public_2", "upstream_2")
|
||||
|
||||
adaptor := &taskPollingFetchAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
|
||||
channelID: {
|
||||
first.GetUpstreamTaskID(),
|
||||
second.GetUpstreamTaskID(),
|
||||
},
|
||||
}, map[string]*model.Task{
|
||||
first.GetUpstreamTaskID(): first,
|
||||
second.GetUpstreamTaskID(): second,
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
assert.Equal(t, 1, adaptor.fetchCount())
|
||||
}
|
||||
|
||||
func TestDispatchPlatformUpdateUsesFetchMode(t *testing.T) {
|
||||
truncate(t)
|
||||
const channelID = 109
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
task := seedPollingTask(t, channelID, "task_batch", "upstream_batch")
|
||||
taskChannels := map[int][]string{channelID: {task.GetUpstreamTaskID()}}
|
||||
tasks := map[string]*model.Task{task.GetUpstreamTaskID(): task}
|
||||
|
||||
batch := &batchPollingAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return batch }
|
||||
DispatchPlatformUpdate(context.Background(), "batch-plugin", taskChannels, tasks)
|
||||
assert.Equal(t, 1, batch.batchCalls)
|
||||
assert.Equal(t, 0, batch.fetchCount())
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.Equal(t, "https://example.com/result", persisted.GetResultURL())
|
||||
|
||||
perTask := &taskPollingFetchAdaptor{}
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return perTask }
|
||||
DispatchPlatformUpdate(context.Background(), "per-task-plugin", taskChannels, tasks)
|
||||
assert.Equal(t, 1, perTask.fetchCount())
|
||||
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return nil }
|
||||
assert.NotPanics(t, func() { DispatchPlatformUpdate(context.Background(), "missing-plugin", taskChannels, tasks) })
|
||||
GetTaskAdaptorFunc = previousFactory
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksSettlesTieredUsageForTerminalStates(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
status model.TaskStatus
|
||||
units float64
|
||||
actualQuota int
|
||||
}{
|
||||
{name: "success with usage", status: model.TaskStatusSuccess, units: 3, actualQuota: 3_000},
|
||||
{name: "failure with usage", status: model.TaskStatusFailure, units: 3, actualQuota: 0},
|
||||
{name: "success with zero usage", status: model.TaskStatusSuccess, units: 0, actualQuota: 0},
|
||||
{name: "failure with zero usage", status: model.TaskStatusFailure, units: 0, actualQuota: 0},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 41, 41, 141
|
||||
const initialQuota, preConsumedQuota = 10_000, 5_000
|
||||
const tokenRemain = 8_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-tiered", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
expression := `tier("actual", u("units"))`
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_tiered_" + string(testCase.status)
|
||||
task.Platform = "batch-plugin"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_tiered_" + string(testCase.status)
|
||||
task.SetData(map[string]any{"provider_payload": "must-be-preserved"})
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1_000,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
reason := ""
|
||||
if testCase.status == model.TaskStatusFailure {
|
||||
reason = "upstream failed"
|
||||
}
|
||||
result := &BatchTaskResult{TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamID,
|
||||
Status: string(testCase.status),
|
||||
Reason: reason,
|
||||
UsageFacts: map[string]any{"units": testCase.units},
|
||||
}}
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{upstreamID: result}}
|
||||
taskIDs := []string{upstreamID}
|
||||
taskMap := map[string]*model.Task{upstreamID: task}
|
||||
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: taskIDs}, taskMap))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.Equal(t, testCase.status, persisted.Status)
|
||||
assert.Equal(t, testCase.actualQuota, persisted.Quota)
|
||||
var persistedData map[string]any
|
||||
require.NoError(t, common.Unmarshal(persisted.Data, &persistedData))
|
||||
assert.Equal(t, "must-be-preserved", persistedData["provider_payload"])
|
||||
assert.Equal(t, initialQuota+(preConsumedQuota-testCase.actualQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+(preConsumedQuota-testCase.actualQuota), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
if testCase.status == model.TaskStatusFailure {
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
}
|
||||
|
||||
// A duplicate terminal response must not settle the same task twice.
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: taskIDs}, taskMap))
|
||||
assert.Equal(t, initialQuota+(preConsumedQuota-testCase.actualQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksRefundsFailedTieredTask(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 43, 43, 143
|
||||
const initialQuota, preConsumedQuota, tokenRemain = 10_000, 5_000, 8_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-tiered-refund", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
expression := `tier("actual", u("units"))`
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_tiered_refund"
|
||||
task.Platform = "batch-plugin"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_tiered_refund"
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1_000,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: map[string]any{"units": float64(5)},
|
||||
EstimatedTier: "actual",
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamID: {TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamID,
|
||||
Status: model.TaskStatusFailure,
|
||||
Reason: "upstream failed",
|
||||
UsageFacts: map[string]any{"units": float64(5)},
|
||||
}},
|
||||
}}
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.EqualValues(t, model.TaskStatusFailure, persisted.Status)
|
||||
assert.Zero(t, persisted.Quota)
|
||||
assert.Equal(t, initialQuota+preConsumedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumedQuota, getTokenRemainQuota(t, tokenID))
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
assert.Equal(t, preConsumedQuota, log.Quota)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, "actual", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, map[string]any{"units": float64(5)}, facts)
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksRefundsFailedTaskWithoutUsageSettlement(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 42, 42, 142
|
||||
const initialQuota, preConsumedQuota, tokenRemain = 10_000, 4_000, 7_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-refund", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_refund"
|
||||
task.Platform = "batch-plugin"
|
||||
task.Properties.OriginModelName = "missing-batch-token-price"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_refund"
|
||||
task.PrivateData.BillingContext.OriginModelName = "missing-batch-token-price"
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamID: {TaskInfo: relaycommon.TaskInfo{TaskID: upstreamID, Status: model.TaskStatusFailure, Reason: "upstream failed", TotalTokens: 123}},
|
||||
}}
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
|
||||
|
||||
assert.Equal(t, initialQuota+preConsumedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumedQuota, getTokenRemainQuota(t, tokenID))
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksCanSkipPollingSleepPerChannel(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const channelID = 102
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
first := seedPollingTask(t, channelID, "task_public_3", "upstream_3")
|
||||
second := seedPollingTask(t, channelID, "task_public_4", "upstream_4")
|
||||
|
||||
adaptor := &taskPollingFetchAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
|
||||
channelID: {
|
||||
first.GetUpstreamTaskID(),
|
||||
second.GetUpstreamTaskID(),
|
||||
},
|
||||
}, map[string]*model.Task{
|
||||
first.GetUpstreamTaskID(): first,
|
||||
second.GetUpstreamTaskID(): second,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, adaptor.fetchCount())
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksDefaultSleepDoesNotBlockOtherChannels(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const firstChannelID = 201
|
||||
const secondChannelID = 202
|
||||
seedTaskPollingChannel(t, firstChannelID, false)
|
||||
seedTaskPollingChannel(t, secondChannelID, false)
|
||||
firstChannelFirst := seedPollingTask(t, firstChannelID, "task_public_5", "upstream_a_1")
|
||||
firstChannelSecond := seedPollingTask(t, firstChannelID, "task_public_6", "upstream_a_2")
|
||||
secondChannelFirst := seedPollingTask(t, secondChannelID, "task_public_7", "upstream_b_1")
|
||||
secondChannelSecond := seedPollingTask(t, secondChannelID, "task_public_8", "upstream_b_2")
|
||||
|
||||
adaptor := &taskPollingFetchAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
|
||||
firstChannelID: {
|
||||
firstChannelFirst.GetUpstreamTaskID(),
|
||||
firstChannelSecond.GetUpstreamTaskID(),
|
||||
},
|
||||
secondChannelID: {
|
||||
secondChannelFirst.GetUpstreamTaskID(),
|
||||
secondChannelSecond.GetUpstreamTaskID(),
|
||||
},
|
||||
}, map[string]*model.Task{
|
||||
firstChannelFirst.GetUpstreamTaskID(): firstChannelFirst,
|
||||
firstChannelSecond.GetUpstreamTaskID(): firstChannelSecond,
|
||||
secondChannelFirst.GetUpstreamTaskID(): secondChannelFirst,
|
||||
secondChannelSecond.GetUpstreamTaskID(): secondChannelSecond,
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
assert.ElementsMatch(t, []string{"upstream_a_1", "upstream_b_1"}, adaptor.fetchedTaskIDs())
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksSlowChannelDoesNotBlockOtherChannels(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const slowChannelID = 251
|
||||
const fastChannelID = 252
|
||||
seedTaskPollingChannel(t, slowChannelID, false)
|
||||
seedTaskPollingChannel(t, fastChannelID, true)
|
||||
slowTask := seedPollingTask(t, slowChannelID, "task_public_slow", "upstream_slow_1")
|
||||
fastFirst := seedPollingTask(t, fastChannelID, "task_public_fast_1", "upstream_fast_parallel_1")
|
||||
fastSecond := seedPollingTask(t, fastChannelID, "task_public_fast_2", "upstream_fast_parallel_2")
|
||||
|
||||
adaptor := &taskPollingFetchAdaptor{
|
||||
fetched: make(chan string, 4),
|
||||
blockTaskID: slowTask.GetUpstreamTaskID(),
|
||||
blockStarted: make(chan struct{}),
|
||||
releaseBlock: make(chan struct{}),
|
||||
}
|
||||
var releaseOnce sync.Once
|
||||
releaseBlockedTask := func() {
|
||||
releaseOnce.Do(func() {
|
||||
close(adaptor.releaseBlock)
|
||||
})
|
||||
}
|
||||
t.Cleanup(releaseBlockedTask)
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
gopool.Go(func() {
|
||||
errCh <- UpdateVideoTasks(context.Background(), constant.TaskPlatform("kling"), map[int][]string{
|
||||
slowChannelID: {
|
||||
slowTask.GetUpstreamTaskID(),
|
||||
},
|
||||
fastChannelID: {
|
||||
fastFirst.GetUpstreamTaskID(),
|
||||
fastSecond.GetUpstreamTaskID(),
|
||||
},
|
||||
}, map[string]*model.Task{
|
||||
slowTask.GetUpstreamTaskID(): slowTask,
|
||||
fastFirst.GetUpstreamTaskID(): fastFirst,
|
||||
fastSecond.GetUpstreamTaskID(): fastSecond,
|
||||
})
|
||||
})
|
||||
|
||||
select {
|
||||
case <-adaptor.blockStarted:
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
t.Fatal("slow channel did not start blocking")
|
||||
}
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
fetchedTaskIDs := adaptor.fetchedTaskIDs()
|
||||
return len(fetchedTaskIDs) == 2 &&
|
||||
fetchedTaskIDs[0] == fastFirst.GetUpstreamTaskID() &&
|
||||
fetchedTaskIDs[1] == fastSecond.GetUpstreamTaskID()
|
||||
}, 500*time.Millisecond, 10*time.Millisecond)
|
||||
|
||||
releaseBlockedTask()
|
||||
require.NoError(t, <-errCh)
|
||||
assert.ElementsMatch(t, []string{
|
||||
slowTask.GetUpstreamTaskID(),
|
||||
fastFirst.GetUpstreamTaskID(),
|
||||
fastSecond.GetUpstreamTaskID(),
|
||||
}, adaptor.fetchedTaskIDs())
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksMixedChannelSleepSettings(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const sleepyChannelID = 301
|
||||
const fastChannelID = 302
|
||||
seedTaskPollingChannel(t, sleepyChannelID, false)
|
||||
seedTaskPollingChannel(t, fastChannelID, true)
|
||||
sleepyFirst := seedPollingTask(t, sleepyChannelID, "task_public_9", "upstream_sleepy_1")
|
||||
sleepySecond := seedPollingTask(t, sleepyChannelID, "task_public_10", "upstream_sleepy_2")
|
||||
fastFirst := seedPollingTask(t, fastChannelID, "task_public_11", "upstream_fast_1")
|
||||
fastSecond := seedPollingTask(t, fastChannelID, "task_public_12", "upstream_fast_2")
|
||||
|
||||
adaptor := &taskPollingFetchAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
err := UpdateVideoTasks(ctx, constant.TaskPlatform("kling"), map[int][]string{
|
||||
sleepyChannelID: {
|
||||
sleepyFirst.GetUpstreamTaskID(),
|
||||
sleepySecond.GetUpstreamTaskID(),
|
||||
},
|
||||
fastChannelID: {
|
||||
fastFirst.GetUpstreamTaskID(),
|
||||
fastSecond.GetUpstreamTaskID(),
|
||||
},
|
||||
}, map[string]*model.Task{
|
||||
sleepyFirst.GetUpstreamTaskID(): sleepyFirst,
|
||||
sleepySecond.GetUpstreamTaskID(): sleepySecond,
|
||||
fastFirst.GetUpstreamTaskID(): fastFirst,
|
||||
fastSecond.GetUpstreamTaskID(): fastSecond,
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
assert.ElementsMatch(t, []string{"upstream_sleepy_1", "upstream_fast_1", "upstream_fast_2"}, adaptor.fetchedTaskIDs())
|
||||
}
|
||||
|
||||
func TestUpdateSunoTasksStalePollsRefundExactlyOnce(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 401, 401, 401
|
||||
const initialUserQuota, initialTokenQuota, taskQuota = 10_000, 6_000, 2_500
|
||||
const publicTaskID, upstreamTaskID = "suno_public_refund_once", "suno_upstream_refund_once"
|
||||
|
||||
seedUser(t, userID, initialUserQuota)
|
||||
seedToken(t, tokenID, userID, "sk-suno-refund-once", initialTokenQuota)
|
||||
baseURL := "https://suno.invalid"
|
||||
require.NoError(t, model.DB.Create(&model.Channel{
|
||||
Id: channelID,
|
||||
Type: constant.ChannelTypeSunoAPI,
|
||||
Name: "suno_refund_once",
|
||||
Key: "sk-suno-channel",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
BaseURL: &baseURL,
|
||||
}).Error)
|
||||
|
||||
task := makeTask(userID, channelID, taskQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = publicTaskID
|
||||
task.Platform = constant.TaskPlatformSuno
|
||||
task.Status = model.TaskStatusInProgress
|
||||
task.Progress = "50%"
|
||||
task.SubmitTime = time.Now().Unix()
|
||||
task.PrivateData.UpstreamTaskID = upstreamTaskID
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
var firstPollTask model.Task
|
||||
var staleSecondPollTask model.Task
|
||||
require.NoError(t, model.DB.First(&firstPollTask, task.ID).Error)
|
||||
require.NoError(t, model.DB.First(&staleSecondPollTask, task.ID).Error)
|
||||
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamTaskID: {TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamTaskID,
|
||||
Status: model.TaskStatusFailure,
|
||||
Reason: "upstream failed",
|
||||
}},
|
||||
}}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
require.NoError(t, updateBatchTasks(context.Background(), adaptor, channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
upstreamTaskID: &firstPollTask,
|
||||
}))
|
||||
require.NoError(t, updateBatchTasks(context.Background(), adaptor, channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
upstreamTaskID: &staleSecondPollTask,
|
||||
}))
|
||||
|
||||
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)
|
||||
assert.Equal(t, initialUserQuota+taskQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota+taskQuota, getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
}
|
||||
|
||||
func TestRunTaskPollingOnceDoesNotRefundHistoricalFailedTask(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, initialQuota, taskQuota = 402, 10_000, 1_200
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
task := makeTask(userID, 0, taskQuota, 0, BillingSourceWallet, 0)
|
||||
task.TaskID = "historical_failed_already_refunded"
|
||||
task.Status = model.TaskStatusFailure
|
||||
task.Progress = "100%"
|
||||
task.SubmitTime = time.Now().Add(-90 * 24 * time.Hour).Unix()
|
||||
task.UpdatedAt = time.Now().Add(-time.Minute).Unix()
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor {
|
||||
return &taskPollingFetchAdaptor{}
|
||||
}
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
summary := RunTaskPollingOnce(context.Background(), nil)
|
||||
|
||||
assert.Zero(t, summary.UnfinishedTasks)
|
||||
assert.Equal(t, initialQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, taskQuota, getTaskQuota(t, task.ID))
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
}
|
||||
|
||||
func TestSweepTimedOutTasksHonorsRefundRolloutBoundary(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const (
|
||||
userID = 403
|
||||
initialQuota = 10_000
|
||||
legacyTaskQuota = 1_800
|
||||
modernTaskQuota = 1_200
|
||||
)
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
legacyTask := makeTask(userID, 0, legacyTaskQuota, 0, BillingSourceWallet, 0)
|
||||
legacyTask.TaskID = "legacy_timeout_without_refund"
|
||||
legacyTask.Progress = "50%"
|
||||
legacyTask.SubmitTime = 1771718399 // 2026-02-21 23:59:59 UTC
|
||||
require.NoError(t, model.DB.Create(legacyTask).Error)
|
||||
|
||||
modernTask := makeTask(userID, 0, modernTaskQuota, 0, BillingSourceWallet, 0)
|
||||
modernTask.TaskID = "modern_timeout_with_refund"
|
||||
modernTask.Progress = "50%"
|
||||
modernTask.SubmitTime = 1771718400 // 2026-02-22 00:00:00 UTC
|
||||
require.NoError(t, model.DB.Create(modernTask).Error)
|
||||
|
||||
previousTimeout := constant.TaskTimeoutMinutes
|
||||
constant.TaskTimeoutMinutes = 1
|
||||
t.Cleanup(func() { constant.TaskTimeoutMinutes = previousTimeout })
|
||||
|
||||
sweepTimedOutTasks(context.Background())
|
||||
|
||||
var reloadedLegacy model.Task
|
||||
var reloadedModern model.Task
|
||||
require.NoError(t, model.DB.First(&reloadedLegacy, legacyTask.ID).Error)
|
||||
require.NoError(t, model.DB.First(&reloadedModern, modernTask.ID).Error)
|
||||
assert.EqualValues(t, model.TaskStatusFailure, reloadedLegacy.Status)
|
||||
assert.EqualValues(t, model.TaskStatusFailure, reloadedModern.Status)
|
||||
assert.Zero(t, reloadedLegacy.Quota)
|
||||
assert.Zero(t, reloadedModern.Quota)
|
||||
assert.Contains(t, reloadedLegacy.FailReason, "旧系统遗留任务")
|
||||
assert.Contains(t, reloadedModern.FailReason, "任务超时")
|
||||
assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
}
|
||||
|
||||
type scriptedPollingAdaptor struct {
|
||||
statusCode int
|
||||
body []byte
|
||||
fetchErr error
|
||||
parse *relaycommon.TaskInfo
|
||||
parseErr error
|
||||
}
|
||||
|
||||
func (a *scriptedPollingAdaptor) Init(*relaycommon.RelayInfo) {}
|
||||
func (a *scriptedPollingAdaptor) FetchTask(string, string, *model.Task, string) (*http.Response, error) {
|
||||
if a.fetchErr != nil {
|
||||
return nil, a.fetchErr
|
||||
}
|
||||
code := a.statusCode
|
||||
if code == 0 {
|
||||
code = http.StatusOK
|
||||
}
|
||||
body := a.body
|
||||
if body == nil {
|
||||
body = []byte(`{}`)
|
||||
}
|
||||
return &http.Response{StatusCode: code, Body: io.NopCloser(bytes.NewReader(body))}, nil
|
||||
}
|
||||
func (a *scriptedPollingAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
|
||||
if a.parseErr != nil {
|
||||
return nil, a.parseErr
|
||||
}
|
||||
if a.parse != nil {
|
||||
return a.parse, nil
|
||||
}
|
||||
return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil
|
||||
}
|
||||
func (a *scriptedPollingAdaptor) AdjustBillingOnComplete(*model.Task, *relaycommon.TaskInfo) int {
|
||||
return 0
|
||||
}
|
||||
|
||||
type scriptedBatchPollingAdaptor struct {
|
||||
scriptedPollingAdaptor
|
||||
results map[string]*BatchTaskResult
|
||||
}
|
||||
|
||||
func (a *scriptedBatchPollingAdaptor) FetchMode() string { return "batch" }
|
||||
func (a *scriptedBatchPollingAdaptor) FetchBatchTasks(string, string, []*model.Task, string) (*http.Response, error) {
|
||||
return a.FetchTask("", "", nil, "")
|
||||
}
|
||||
func (a *scriptedBatchPollingAdaptor) ParseBatchResult([]*model.Task, *http.Response, []byte) (map[string]*BatchTaskResult, error) {
|
||||
if a.parseErr != nil {
|
||||
return nil, a.parseErr
|
||||
}
|
||||
return a.results, nil
|
||||
}
|
||||
|
||||
func TestUpdateVideoSingleTaskPollClassification(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
fetchErr error
|
||||
parse *relaycommon.TaskInfo
|
||||
parseErr error
|
||||
priorFailures int
|
||||
priorState string
|
||||
maxFailures int
|
||||
wantStatus model.TaskStatus
|
||||
wantFailures int
|
||||
wantRefund bool
|
||||
wantReason string
|
||||
wantState string
|
||||
wantUnchanged bool
|
||||
}{
|
||||
{
|
||||
name: "404 fails immediately and refunds",
|
||||
statusCode: http.StatusNotFound,
|
||||
wantStatus: model.TaskStatusFailure,
|
||||
wantRefund: true,
|
||||
wantReason: "upstream task not found (HTTP 404)",
|
||||
wantUnchanged: false,
|
||||
},
|
||||
{
|
||||
name: "401 increments without changing status",
|
||||
statusCode: http.StatusUnauthorized,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 1,
|
||||
wantUnchanged: true,
|
||||
},
|
||||
{
|
||||
name: "429 reaches threshold and refunds",
|
||||
statusCode: http.StatusTooManyRequests,
|
||||
priorFailures: 2,
|
||||
maxFailures: 3,
|
||||
wantStatus: model.TaskStatusFailure,
|
||||
wantFailures: 3,
|
||||
wantRefund: true,
|
||||
wantReason: "poll failed: transient (HTTP 429)",
|
||||
},
|
||||
{
|
||||
name: "UNKNOWN increments",
|
||||
statusCode: http.StatusOK,
|
||||
parse: &relaycommon.TaskInfo{Status: model.TaskStatusUnknown, Reason: "weird"},
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 1,
|
||||
wantUnchanged: true,
|
||||
},
|
||||
{
|
||||
name: "valid 2xx resets the failure counter",
|
||||
statusCode: http.StatusOK,
|
||||
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
|
||||
priorFailures: 5,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 0,
|
||||
},
|
||||
{
|
||||
name: "omit state preserves previous plugin state",
|
||||
statusCode: http.StatusOK,
|
||||
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
|
||||
priorState: `{"req_key":"keep"}`,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantState: `{"req_key":"keep"}`,
|
||||
},
|
||||
{
|
||||
name: "returned state replaces plugin state",
|
||||
statusCode: http.StatusOK,
|
||||
parse: &relaycommon.TaskInfo{
|
||||
Status: model.TaskStatusInProgress,
|
||||
PluginState: []byte(`{"req_key":"new"}`),
|
||||
},
|
||||
priorState: `{"req_key":"old"}`,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantState: `{"req_key":"new"}`,
|
||||
},
|
||||
{
|
||||
name: "other 4xx non-terminal is unrecognized",
|
||||
statusCode: http.StatusBadRequest,
|
||||
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 1,
|
||||
wantUnchanged: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID, tokenID, channelID = 510, 510, 510
|
||||
const initialQuota, preConsumed, tokenRemain = 10_000, 4_000, 7_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-poll-class", tokenRemain)
|
||||
ch := &model.Channel{Id: channelID, Type: constant.ChannelTypeKling, Name: "poll", Key: "sk-test", Status: common.ChannelStatusEnabled}
|
||||
|
||||
if testCase.maxFailures > 0 {
|
||||
previous := constant.TaskPollMaxFailures
|
||||
constant.TaskPollMaxFailures = testCase.maxFailures
|
||||
t.Cleanup(func() { constant.TaskPollMaxFailures = previous })
|
||||
}
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_poll_class"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_poll_class"
|
||||
task.PrivateData.PollFailures = testCase.priorFailures
|
||||
if testCase.priorState != "" {
|
||||
task.PrivateData.PluginState = []byte(testCase.priorState)
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
adaptor := &scriptedPollingAdaptor{statusCode: testCase.statusCode, fetchErr: testCase.fetchErr, parse: testCase.parse, parseErr: testCase.parseErr}
|
||||
require.NoError(t, updateVideoSingleTask(context.Background(), adaptor, ch, task.GetUpstreamTaskID(), map[string]*model.Task{
|
||||
task.GetUpstreamTaskID(): task,
|
||||
}))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.EqualValues(t, testCase.wantStatus, persisted.Status)
|
||||
assert.Equal(t, testCase.wantFailures, persisted.PrivateData.PollFailures)
|
||||
if testCase.wantUnchanged {
|
||||
assert.Empty(t, persisted.FailReason)
|
||||
}
|
||||
if testCase.wantReason != "" {
|
||||
assert.Contains(t, persisted.FailReason, testCase.wantReason)
|
||||
}
|
||||
if testCase.wantState != "" {
|
||||
assert.JSONEq(t, testCase.wantState, string(persisted.PrivateData.PluginState))
|
||||
}
|
||||
if testCase.wantRefund {
|
||||
assert.Equal(t, initialQuota+preConsumed, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
||||
assert.Zero(t, persisted.Quota)
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
} else {
|
||||
assert.Equal(t, initialQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksPollClassification(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
resultStatus model.TaskStatus
|
||||
wantStatus model.TaskStatus
|
||||
wantFailures int
|
||||
wantRefund bool
|
||||
wantReason string
|
||||
}{
|
||||
{
|
||||
name: "404 fails the batch and refunds",
|
||||
statusCode: http.StatusNotFound,
|
||||
wantStatus: model.TaskStatusFailure,
|
||||
wantRefund: true,
|
||||
wantReason: "upstream task not found (HTTP 404)",
|
||||
},
|
||||
{
|
||||
name: "401 increments every task",
|
||||
statusCode: http.StatusUnauthorized,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 1,
|
||||
},
|
||||
{
|
||||
name: "UNKNOWN increments",
|
||||
statusCode: http.StatusOK,
|
||||
resultStatus: model.TaskStatusUnknown,
|
||||
wantStatus: model.TaskStatusInProgress,
|
||||
wantFailures: 1,
|
||||
},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID, tokenID, channelID = 610, 610, 610
|
||||
const initialQuota, preConsumed, tokenRemain = 10_000, 4_000, 7_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-class", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_class"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_class"
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
|
||||
adaptor := &scriptedBatchPollingAdaptor{
|
||||
scriptedPollingAdaptor: scriptedPollingAdaptor{statusCode: testCase.statusCode},
|
||||
}
|
||||
if testCase.resultStatus != "" {
|
||||
adaptor.results = map[string]*BatchTaskResult{
|
||||
upstreamID: {TaskInfo: relaycommon.TaskInfo{TaskID: upstreamID, Status: string(testCase.resultStatus), Reason: "weird"}},
|
||||
}
|
||||
}
|
||||
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.EqualValues(t, testCase.wantStatus, persisted.Status)
|
||||
assert.Equal(t, testCase.wantFailures, persisted.PrivateData.PollFailures)
|
||||
if testCase.wantReason != "" {
|
||||
assert.Contains(t, persisted.FailReason, testCase.wantReason)
|
||||
}
|
||||
if testCase.wantRefund {
|
||||
assert.Equal(t, initialQuota+preConsumed, getUserQuota(t, userID))
|
||||
assert.Zero(t, persisted.Quota)
|
||||
} else {
|
||||
assert.Equal(t, initialQuota, getUserQuota(t, userID))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user