cneicy
a5f7b8e586
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
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
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
451 lines
16 KiB
Go
451 lines
16 KiB
Go
package middleware
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
appI18n "github.com/QuantumNous/new-api/i18n"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func setupOriginTaskDB(t *testing.T) {
|
|
t.Helper()
|
|
previousDB := model.DB
|
|
previousType := common.MainDatabaseType()
|
|
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}))
|
|
model.DB = database
|
|
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
|
t.Cleanup(func() {
|
|
model.DB = previousDB
|
|
common.SetMainDatabaseType(previousType)
|
|
})
|
|
}
|
|
|
|
func insertOriginTaskChannel(t *testing.T, status int) *model.Channel {
|
|
t.Helper()
|
|
channel := &model.Channel{
|
|
Name: "origin-channel",
|
|
Key: "sk-origin",
|
|
Status: status,
|
|
Type: constant.ChannelTypeDoubaoVideo,
|
|
}
|
|
require.NoError(t, model.DB.Create(channel).Error)
|
|
return channel
|
|
}
|
|
|
|
func insertOriginOwnedTask(t *testing.T, taskID string, userID, channelID int, platform constant.TaskPlatform) *model.Task {
|
|
t.Helper()
|
|
task := &model.Task{
|
|
TaskID: taskID,
|
|
UserId: userID,
|
|
ChannelId: channelID,
|
|
Platform: platform,
|
|
Action: "text_to_video",
|
|
Status: model.TaskStatusSuccess,
|
|
PrivateData: model.TaskPrivateData{
|
|
UpstreamTaskID: "upstream-" + taskID,
|
|
},
|
|
}
|
|
data, err := common.Marshal(map[string]any{"id": "upstream-" + taskID})
|
|
require.NoError(t, err)
|
|
task.Data = data
|
|
require.NoError(t, model.DB.Create(task).Error)
|
|
return task
|
|
}
|
|
|
|
func originTaskTestContext(userID int) *gin.Context {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserId, userID)
|
|
return c
|
|
}
|
|
|
|
func resolvedOriginPin(c *gin.Context) (int, bool) {
|
|
pin, found, _ := service.GetChannelConstraints(c).ResolvedPin()
|
|
if !found {
|
|
return 0, false
|
|
}
|
|
return pin.ChannelId, true
|
|
}
|
|
|
|
func TestApplyOriginTaskIntent(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
enabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
otherEnabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
disabled := insertOriginTaskChannel(t, common.ChannelStatusManuallyDisabled)
|
|
insertOriginOwnedTask(t, "task-own", 7, enabled.Id, "origin-plugin")
|
|
insertOriginOwnedTask(t, "task-own-b", 7, enabled.Id, "origin-plugin")
|
|
insertOriginOwnedTask(t, "task-other-channel", 7, otherEnabled.Id, "origin-plugin")
|
|
insertOriginOwnedTask(t, "task-foreign", 8, enabled.Id, "origin-plugin")
|
|
insertOriginOwnedTask(t, "task-wrong-platform", 7, enabled.Id, "other-plugin")
|
|
insertOriginOwnedTask(t, "task-legacy", 7, enabled.Id, constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeDoubaoVideo)))
|
|
insertOriginOwnedTask(t, "task-disabled", 7, disabled.Id, "origin-plugin")
|
|
|
|
tooMany := make([]any, 17)
|
|
for i := range tooMany {
|
|
tooMany[i] = fmt.Sprintf("task-%d", i)
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
userID int
|
|
intent map[string]any
|
|
channelType int
|
|
wantCode string
|
|
wantPinned int
|
|
wantIDs []string
|
|
}{
|
|
{
|
|
name: "valid single origin id pins channel",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-own"}},
|
|
wantPinned: enabled.Id,
|
|
wantIDs: []string{"task-own"},
|
|
},
|
|
{
|
|
name: "dedupes preserving first order",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-own", " task-own ", "task-own-b"}},
|
|
wantPinned: enabled.Id,
|
|
wantIDs: []string{"task-own", "task-own-b"},
|
|
},
|
|
{
|
|
name: "legacy platform matches channelType",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-legacy"}},
|
|
channelType: constant.ChannelTypeDoubaoVideo,
|
|
wantPinned: enabled.Id,
|
|
wantIDs: []string{"task-legacy"},
|
|
},
|
|
{
|
|
name: "unknown id",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-missing"}},
|
|
wantCode: "origin_task_not_found",
|
|
},
|
|
{
|
|
name: "other user's task is not found",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-foreign"}},
|
|
wantCode: "origin_task_not_found",
|
|
},
|
|
{
|
|
name: "platform mismatch",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-wrong-platform"}},
|
|
wantCode: "origin_task_platform_mismatch",
|
|
},
|
|
{
|
|
name: "two ids on different channels",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-own", "task-other-channel"}},
|
|
wantCode: "origin_task_channel_conflict",
|
|
},
|
|
{
|
|
name: "disabled channel",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-disabled"}},
|
|
wantCode: "origin_task_channel_disabled",
|
|
},
|
|
{
|
|
name: "more than 16 ids",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": tooMany},
|
|
wantCode: "invalid_origin_task_ids",
|
|
},
|
|
{
|
|
name: "non-array originTaskIds",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": "task-own"},
|
|
wantCode: "invalid_origin_task_ids",
|
|
},
|
|
{
|
|
name: "empty string entry",
|
|
userID: 7,
|
|
intent: map[string]any{"originTaskIds": []any{"task-own", " "}},
|
|
wantCode: "invalid_origin_task_ids",
|
|
},
|
|
}
|
|
|
|
for _, testCase := range tests {
|
|
t.Run(testCase.name, func(t *testing.T) {
|
|
c := originTaskTestContext(testCase.userID)
|
|
intentErr := applyOriginTaskIntent(c, testCase.intent, jsplugin.Meta{Key: "origin-plugin", ChannelTypes: []int{testCase.channelType}})
|
|
if testCase.wantCode != "" {
|
|
require.NotNil(t, intentErr)
|
|
assert.Equal(t, testCase.wantCode, intentErr.Code)
|
|
assert.Equal(t, http.StatusBadRequest, intentErr.StatusCode)
|
|
_, pinned := resolvedOriginPin(c)
|
|
assert.False(t, pinned)
|
|
return
|
|
}
|
|
require.Nil(t, intentErr)
|
|
pinnedID, ok := resolvedOriginPin(c)
|
|
require.True(t, ok)
|
|
assert.Equal(t, testCase.wantPinned, pinnedID)
|
|
tasks, ok := common.GetContextKeyType[[]*model.Task](c, constant.ContextKeyOriginTasks)
|
|
require.True(t, ok)
|
|
require.Len(t, tasks, len(testCase.wantIDs))
|
|
for i, wantID := range testCase.wantIDs {
|
|
assert.Equal(t, wantID, tasks[i].TaskID)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestApplyOriginTaskIntentAbsentAndEmptyAreNoop(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
c := originTaskTestContext(7)
|
|
require.Nil(t, applyOriginTaskIntent(c, map[string]any{}, jsplugin.Meta{Key: "origin-plugin"}))
|
|
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{}}, jsplugin.Meta{Key: "origin-plugin"}))
|
|
_, pinned := resolvedOriginPin(c)
|
|
assert.False(t, pinned)
|
|
}
|
|
|
|
func TestPrepareTaskPluginRoutePinsOriginTaskChannel(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
insertOriginOwnedTask(t, "task-route", 7, channel.Id, "origin-route")
|
|
plugin := compileTaskRoutePlugin(t, `
|
|
export const meta = {
|
|
apiVersion: 1, key: "origin-route", name: "Origin", version: "1.0.0",
|
|
author: {name: "Test"},
|
|
models: ["resolved-model"], fetchMode: "per_task",
|
|
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
|
|
};
|
|
export const native = {
|
|
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["task-route"], requestBody: {prompt: "ok"}}; },
|
|
jobCreated: function(ctx, task) { return task; },
|
|
};
|
|
export function buildSubmitRequest() { return {url: "https://example.com"}; }
|
|
export function parseSubmitResponse() { return {taskId: "one"}; }
|
|
export function buildQueryRequest() { return {url: "https://example.com"}; }
|
|
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
|
`)
|
|
reached := false
|
|
router := gin.New()
|
|
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
|
|
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
|
c.Next()
|
|
}, PrepareTaskPluginRoute(), func(c *gin.Context) {
|
|
reached = true
|
|
pinnedID, ok := resolvedOriginPin(c)
|
|
require.True(t, ok)
|
|
assert.Equal(t, channel.Id, pinnedID)
|
|
c.Status(http.StatusNoContent)
|
|
})
|
|
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, request)
|
|
assert.True(t, reached)
|
|
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
|
}
|
|
|
|
func TestPrepareTaskPluginRouteRejectsUnknownOriginTask(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
plugin := compileTaskRoutePlugin(t, `
|
|
export const meta = {
|
|
apiVersion: 1, key: "origin-route-missing", name: "Origin", version: "1.0.0",
|
|
author: {name: "Test"},
|
|
models: ["resolved-model"], fetchMode: "per_task",
|
|
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
|
|
};
|
|
export const native = {
|
|
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["missing"], requestBody: {}}; },
|
|
jobCreated: function(ctx, task) { return task; },
|
|
};
|
|
export function buildSubmitRequest() { return {url: "https://example.com"}; }
|
|
export function parseSubmitResponse() { return {taskId: "one"}; }
|
|
export function buildQueryRequest() { return {url: "https://example.com"}; }
|
|
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
|
`)
|
|
reached := false
|
|
router := gin.New()
|
|
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
|
|
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
|
c.Next()
|
|
}, PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true })
|
|
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, request)
|
|
assert.False(t, reached)
|
|
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
}
|
|
|
|
func TestPrepareTaskPluginEndpointPinsOriginTaskChannel(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
insertOriginOwnedTask(t, "task-endpoint", 7, channel.Id, "origin-endpoint")
|
|
const key = "origin-endpoint"
|
|
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
|
|
key,
|
|
"1.0.0",
|
|
`["claimed-model"]`,
|
|
"/v1/responses",
|
|
`return {model: ctx.model, originTaskIds: ["task-endpoint"], requestBody: {prompt: "ok"}};`,
|
|
), jsplugin.Options{})
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
|
|
|
|
reached := false
|
|
router := gin.New()
|
|
router.POST("/v1/responses", func(c *gin.Context) {
|
|
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
|
c.Next()
|
|
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) {
|
|
reached = true
|
|
pinnedID, ok := resolvedOriginPin(c)
|
|
require.True(t, ok)
|
|
assert.Equal(t, channel.Id, pinnedID)
|
|
c.Status(http.StatusNoContent)
|
|
})
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, request)
|
|
assert.True(t, reached)
|
|
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
|
}
|
|
|
|
func TestPrepareTaskPluginEndpointRejectsUnknownOriginTask(t *testing.T) {
|
|
setupOriginTaskDB(t)
|
|
const key = "origin-endpoint-missing"
|
|
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
|
|
key,
|
|
"1.0.0",
|
|
`["claimed-model"]`,
|
|
"/v1/responses",
|
|
`return {model: ctx.model, originTaskIds: ["missing"], requestBody: {prompt: "ok"}};`,
|
|
), jsplugin.Options{})
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
|
|
|
|
reached := false
|
|
router := gin.New()
|
|
router.POST("/v1/responses", func(c *gin.Context) {
|
|
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
|
c.Next()
|
|
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reached = true })
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
recorder := httptest.NewRecorder()
|
|
router.ServeHTTP(recorder, request)
|
|
assert.False(t, reached)
|
|
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
assert.Contains(t, recorder.Body.String(), "origin_task_not_found")
|
|
}
|
|
|
|
func TestDistributeHonorsOriginTaskChannelPin(t *testing.T) {
|
|
require.NoError(t, appI18n.Init())
|
|
setupOriginTaskDB(t)
|
|
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set("resolved_task_model", "resolved-model")
|
|
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
|
ChannelId: channel.Id,
|
|
Source: dto.PinSourceOriginTask,
|
|
Rank: dto.PinRankOriginTask,
|
|
RetryMode: dto.PinRetrySameChannel,
|
|
})
|
|
|
|
nextCalled := false
|
|
handler := Distribute()
|
|
handler(c)
|
|
if !c.IsAborted() {
|
|
nextCalled = true
|
|
}
|
|
assert.True(t, nextCalled)
|
|
assert.Equal(t, channel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
|
|
}
|
|
|
|
func TestDistributeTokenPinBeatsOriginPin(t *testing.T) {
|
|
require.NoError(t, appI18n.Init())
|
|
setupOriginTaskDB(t)
|
|
tokenChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
originChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
|
|
var warnBuf bytes.Buffer
|
|
previousWriter := gin.DefaultErrorWriter
|
|
gin.DefaultErrorWriter = &warnBuf
|
|
t.Cleanup(func() { gin.DefaultErrorWriter = previousWriter })
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set("resolved_task_model", "resolved-model")
|
|
constraints := service.GetChannelConstraints(c)
|
|
constraints.AddPin(dto.ChannelPin{
|
|
ChannelId: tokenChannel.Id,
|
|
Source: dto.PinSourceToken,
|
|
Rank: dto.PinRankToken,
|
|
RetryMode: dto.PinRetrySingleAttempt,
|
|
})
|
|
constraints.AddPin(dto.ChannelPin{
|
|
ChannelId: originChannel.Id,
|
|
Source: dto.PinSourceOriginTask,
|
|
Rank: dto.PinRankOriginTask,
|
|
RetryMode: dto.PinRetrySameChannel,
|
|
})
|
|
nextCalled := false
|
|
Distribute()(c)
|
|
if !c.IsAborted() {
|
|
nextCalled = true
|
|
}
|
|
assert.True(t, nextCalled)
|
|
assert.Equal(t, tokenChannel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
|
|
warn := warnBuf.String()
|
|
assert.Contains(t, warn, "winning_source=token")
|
|
assert.Contains(t, warn, fmt.Sprintf("winning_channel_id=%d", tokenChannel.Id))
|
|
assert.Contains(t, warn, "overridden_source=origin_task")
|
|
assert.Contains(t, warn, fmt.Sprintf("overridden_channel_id=%d", originChannel.Id))
|
|
}
|
|
|
|
func TestDistributePinViolatingIdentityFilterErrors(t *testing.T) {
|
|
require.NoError(t, appI18n.Init())
|
|
setupOriginTaskDB(t)
|
|
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set("resolved_task_model", "resolved-model")
|
|
c.Set("expected_task_plugin_key", "alpha")
|
|
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
|
ChannelId: channel.Id,
|
|
Source: dto.PinSourceOriginTask,
|
|
Rank: dto.PinRankOriginTask,
|
|
RetryMode: dto.PinRetrySameChannel,
|
|
})
|
|
|
|
Distribute()(c)
|
|
assert.True(t, c.IsAborted())
|
|
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
assert.Contains(t, recorder.Body.String(), string(dto.FilterTaskPluginIdentity))
|
|
}
|