router/video_router_test.go
Rama main · réplica de lectura del repositorio original.
package router
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/plugins"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGetOpenAIVideoRouteRendersJimengTask(t *testing.T) {
gin.SetMode(gin.TestMode)
previousDB := model.DB
previousDatabaseType := common.MainDatabaseType()
previousLogDatabaseType := common.LogDatabaseType()
previousSQLitePath := common.SQLitePath
previousMasterNode := common.IsMasterNode
previousRedisEnabled := common.RedisEnabled
common.SQLitePath = t.TempDir() + "/router-video.db"
common.IsMasterNode = false
common.RedisEnabled = false
t.Setenv("SQL_DSN", "")
require.NoError(t, model.InitDB())
database := model.DB
require.NoError(t, database.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Task{}))
t.Cleanup(func() {
sqlDB, closeErr := database.DB()
require.NoError(t, closeErr)
require.NoError(t, sqlDB.Close())
model.DB = previousDB
common.SetDatabaseTypes(previousDatabaseType, previousLogDatabaseType)
common.SQLitePath = previousSQLitePath
common.IsMasterNode = previousMasterNode
common.RedisEnabled = previousRedisEnabled
})
require.NoError(t, database.Create(&model.User{
Id: 91,
Username: "jimeng-fetch-user",
Role: common.RoleCommonUser,
Status: common.UserStatusEnabled,
Quota: 100,
Group: "default",
AuthVersion: 1,
}).Error)
require.NoError(t, database.Create(&model.Token{
Id: 1,
UserId: 91,
Key: "jimengfetch",
Status: common.TokenStatusEnabled,
Name: "jimeng fetch",
ExpiredTime: -1,
UnlimitedQuota: true,
}).Error)
require.NoError(t, database.Create(&model.Channel{
Id: 17,
Type: constant.ChannelTypeJimeng,
Key: "unused",
Status: common.ChannelStatusEnabled,
Name: "jimeng fetch",
Models: "jimeng_vgfm_t2v_l20",
Group: "default",
}).Error)
task := &model.Task{
CreatedAt: 1710000000,
UpdatedAt: 1710000060,
TaskID: "task_jimeng_public",
Platform: constant.TaskPlatform("jimeng"),
UserId: 91,
Group: "default",
ChannelId: 17,
Status: model.TaskStatusSuccess,
Progress: "100%",
PrivateData: model.TaskPrivateData{
ResultURL: "data:video/mp4;base64,ZGF0YQ==",
},
}
task.SetData(map[string]any{
"code": 10000,
"data": map[string]any{
"status": "done",
"task_id": "jimeng-private-1",
"video_url": "https://cdn.example/video.mp4",
},
"message": "success",
})
require.NoError(t, database.Create(task).Error)
engine := gin.New()
SetVideoRouter(engine)
SetTaskPluginProtocolRouter(engine)
request := httptest.NewRequest(http.MethodGet, "/v1/videos/task_jimeng_public", nil)
request.Header.Set("Authorization", "Bearer sk-jimengfetch")
recorder := httptest.NewRecorder()
engine.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
var response struct {
ID string `json:"id"`
Object string `json:"object"`
Status string `json:"status"`
Progress int `json:"progress"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, "task_jimeng_public", response.ID)
assert.Equal(t, "video", response.Object)
assert.Equal(t, "completed", response.Status)
assert.Equal(t, 100, response.Progress)
assert.Equal(t, int64(1710000000), response.CreatedAt)
assert.Equal(t, int64(1710000060), response.CompletedAt)
assert.NotContains(t, recorder.Body.String(), "cdn.example")
assert.NotContains(t, recorder.Body.String(), "jimeng-private-1")
for _, testCase := range []struct {
name string
authorization string
query string
wantStatus int
}{
{name: "missing credential rejected", wantStatus: http.StatusUnauthorized},
{name: "access rejected", query: "?access=not-a-video-credential", wantStatus: http.StatusUnauthorized},
{name: "bearer accepted", authorization: "Bearer sk-jimengfetch", wantStatus: http.StatusOK},
} {
t.Run(testCase.name, func(t *testing.T) {
request := httptest.NewRequest(
http.MethodGet,
"/v1/videos/task_jimeng_public/content"+testCase.query,
nil,
)
if testCase.authorization != "" {
request.Header.Set("Authorization", testCase.authorization)
}
recorder := httptest.NewRecorder()
engine.ServeHTTP(recorder, request)
assert.Equal(t, testCase.wantStatus, recorder.Code, recorder.Body.String())
if testCase.wantStatus == http.StatusOK {
assert.Equal(t, "data", recorder.Body.String())
}
})
}
}
func TestGetOpenAIVideoRouteHandsNativeXAITaskToPluginQuery(t *testing.T) {
setupRelayRouterTestDB(t)
require.NoError(t, model.DB.AutoMigrate(&model.Channel{}, &model.Task{}))
require.NoError(t, model.DB.Create(&model.User{Id: 92, Username: "xai-fetch-user", Status: common.UserStatusEnabled, Group: "default", Quota: 100}).Error)
require.NoError(t, model.DB.Create(&model.Token{UserId: 92, Key: "xaifetch", Status: common.TokenStatusEnabled, ExpiredTime: -1, UnlimitedQuota: true}).Error)
for taskID, requestPath := range map[string]string{"task_xai_native": "/v1/videos/generations", "task_xai_pinned": "/xai/v1/videos/generations", "task_xai_openai": "/v1/videos"} {
task := &model.Task{
CreatedAt: 1710000000,
UpdatedAt: 1710000060,
TaskID: taskID,
Platform: constant.TaskPlatform("xai"),
UserId: 92,
Group: "default",
Status: model.TaskStatusSuccess,
Progress: "100%",
Properties: model.Properties{OriginModelName: "grok-imagine-video"},
PrivateData: model.TaskPrivateData{
UpstreamTaskID: "xai-private-1",
Execution: &model.TaskExecutionSnapshot{
RequestPath: requestPath,
TaskPlugin: &model.TaskPluginSnapshot{Key: "xai"},
},
},
}
task.SetData(map[string]any{
"request_id": "xai-private-1",
"status": "done",
"model": "grok-imagine-video",
"video": map[string]any{"url": "https://cdn.example/xai.mp4", "duration": 8},
})
require.NoError(t, model.DB.Create(task).Error)
}
engine := gin.New()
SetVideoRouter(engine)
SetTaskPluginProtocolRouter(engine)
source, err := plugins.Source("xai")
require.NoError(t, err)
registry := jsplugin.NewRegistry()
_, err = registry.RegisterFactory(source, jsplugin.Options{Key: "xai"})
require.NoError(t, err)
require.NoError(t, registry.SetGenerationPreparer(newPluginGenerationBuilder(engine.Routes(), nil, productionPluginRouteHandlers).prepare))
require.Empty(t, registry.RoutingErrors())
previousRegistry := defaultPluginDispatcher.registry
defaultPluginDispatcher.registry = registry
t.Cleanup(func() { defaultPluginDispatcher.registry = previousRegistry })
for _, testCase := range []struct {
taskID string
wantStatus string
wantObject any
}{
// Submitted on the xAI route: xAI's own query response.
{taskID: "task_xai_native", wantStatus: "done"},
// The same under /xai.
{taskID: "task_xai_pinned", wantStatus: "done"},
// Submitted through POST /v1/videos: the OpenAI video object.
{taskID: "task_xai_openai", wantStatus: "completed", wantObject: "video"},
} {
t.Run(testCase.taskID, func(t *testing.T) {
request := httptest.NewRequest(http.MethodGet, "/v1/videos/"+testCase.taskID, nil)
request.Header.Set("Authorization", "Bearer sk-xaifetch")
recorder := httptest.NewRecorder()
engine.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
var response map[string]any
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, testCase.wantStatus, response["status"])
assert.Equal(t, testCase.wantObject, response["object"])
assert.NotContains(t, recorder.Body.String(), "xai-private-1")
})
}
}