service/channel_select_test.go
Rama main · réplica de lectura del repositorio original.
package service
import (
"fmt"
"testing"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPinnedTaskPluginChannelTypesUsesPinnedGenerationIndex(t *testing.T) {
registry := jsplugin.NewRegistry()
plugin, err := registry.Register(channelSelectTaskPluginSource("legacy-select", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: plugin,
})
types, keys := pinnedTaskPluginIdentities(c, "legacy-select")
assert.Equal(t, []int{constant.ChannelTypeKling}, types)
assert.Equal(t, []string{"legacy-select"}, keys)
types, keys = pinnedTaskPluginIdentities(c, "another-plugin")
assert.Empty(t, types)
assert.Empty(t, keys)
types, keys = pinnedTaskPluginIdentities(nil, "legacy-select")
assert.Empty(t, types)
assert.Empty(t, keys)
}
func TestPinnedTaskPluginChannelTypesLeavesGenericChannelsKeyed(t *testing.T) {
registry := jsplugin.NewRegistry()
plugin, err := registry.Register(channelSelectTaskPluginSource("generic-select", constant.ChannelTypeTaskPlugin), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: plugin,
})
types, keys := pinnedTaskPluginIdentities(c, "generic-select")
assert.Empty(t, types)
assert.Equal(t, []string{"generic-select"}, keys)
}
func TestPinnedTaskPluginChannelTypesIncludesSharedEndpointProviders(t *testing.T) {
registry := jsplugin.NewRegistry()
_, err := registry.Register(channelSelectEndpointPluginSource("gemini-select", constant.ChannelTypeGemini), jsplugin.Options{})
require.NoError(t, err)
_, err = registry.Register(channelSelectEndpointPluginSource("vertex-select", constant.ChannelTypeVertexAi), jsplugin.Options{})
require.NoError(t, err)
candidates := registry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: candidates[0].Plugin,
})
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
Generation: registry.Generation(),
Plugin: candidates[0].Plugin,
Protocol: candidates[0].Protocol,
Operation: candidates[0].Operation,
Model: "task-model",
Candidates: candidates,
})
AppendTaskPluginIdentityFilter(c, candidates[0].Plugin.Meta.Key)
filters := GetChannelConstraints(c).Filters
require.Len(t, filters, 1)
assert.Equal(t, []int{constant.ChannelTypeGemini, constant.ChannelTypeVertexAi}, filters[0].TaskPluginChannelTypes)
assert.Equal(t, []string{"gemini-select", "vertex-select"}, filters[0].TaskPluginKeys)
}
func channelSelectTaskPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
%s
models: ["task-model"],
fetchMode: "per_task",
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`, key, key, channelSelectChannelTypesField(channelType))
}
func channelSelectEndpointPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
%s
models: ["task-model"],
fetchMode: "per_task",
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
export const protocols = {openai_responses: {
decodeRequest: function(ctx) { return {kind: "submit", model: "task-model", requestBody: ctx.body.value}; },
renderEvents: function() { return {events: [], state: null, done: false}; },
renderFinal: function() { return {output: []}; },
}};
`, key, key, channelSelectChannelTypesField(channelType))
}
func channelSelectChannelTypesField(channelType int) string {
if channelType <= 0 || channelType == constant.ChannelTypeTaskPlugin {
return ""
}
return fmt.Sprintf("channelTypes: [%d],", channelType)
}
func TestPinnedTaskPluginChannelTypesIncludesCompatibleTypes(t *testing.T) {
registry := jsplugin.NewRegistry()
plugin, err := registry.Register(channelSelectCompatiblePluginSource("sora-select", constant.ChannelTypeSora, constant.ChannelTypeOpenAI), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: plugin,
})
types, keys := pinnedTaskPluginIdentities(c, "sora-select")
assert.Equal(t, []int{constant.ChannelTypeSora, constant.ChannelTypeOpenAI}, types)
assert.Equal(t, []string{"sora-select"}, keys)
}
func channelSelectCompatiblePluginSource(key string, channelType, compatibleType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
channelTypes: [%d, %d],
models: ["task-model"],
fetchMode: "per_task",
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`, key, key, channelType, compatibleType)
}
func TestSharedType61IdentityFilterContainsAllCandidateKeys(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"alpha", "beta"} {
_, err := registry.Register(channelSelectEndpointPluginSource(key, 0), jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Generation: generation, Plugin: candidates[0].Plugin, Candidates: candidates})
AppendTaskPluginIdentityFilter(c, "alpha")
filters := GetChannelConstraints(c).Filters
require.Len(t, filters, 1)
assert.Equal(t, "alpha", filters[0].TaskPluginKey)
assert.Equal(t, []string{"alpha", "beta"}, filters[0].TaskPluginKeys)
assert.Empty(t, filters[0].TaskPluginChannelTypes)
}
func TestSharedNativeRouteIdentityFilterContainsAllCandidateKeys(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"route-alpha", "route-beta"} {
_, err := registry.Register(fmt.Sprintf(`
export const meta = {apiVersion: 1, key: %q, name: %q, version: "1.0.0", author: {name: "Test"}, models: ["task-model"], fetchMode: "per_task",
routes: [{method: "POST", path: "/probe/v1/jobs", type: "submit", decode: "decode", render: "render"}]};
export const native = {decode() { return {}; }, render() { return {}; }};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`, key, key), jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupDeclaredRoutes("POST", "/probe/v1/jobs")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Generation: generation, Plugin: candidates[0].Plugin, Route: candidates[0].Route, Candidates: candidates})
AppendTaskPluginIdentityFilter(c, "route-alpha")
filters := GetChannelConstraints(c).Filters
require.Len(t, filters, 1)
assert.Equal(t, "route-alpha", filters[0].TaskPluginKey)
assert.Equal(t, []string{"route-alpha", "route-beta"}, filters[0].TaskPluginKeys)
}