334 lines
10 KiB
Go
334 lines
10 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/model"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"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
|
|
}
|
|
|
|
func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
|
|
|
|
func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) {
|
|
taskID, _ := body["task_id"].(string)
|
|
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 := dto.TaskResponse[model.Task]{
|
|
Code: dto.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([]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 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.TaskActionGenerate,
|
|
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 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())
|
|
}
|