feat: publish v1.0.0-rc.36-custom.1

service/channel_select_test.go

@@ -0,0 +1,152 @@
+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,
+	})
+
+	assert.Equal(t, []int{constant.ChannelTypeKling}, pinnedTaskPluginChannelTypes(c, "legacy-select"))
+	assert.Empty(t, pinnedTaskPluginChannelTypes(c, "another-plugin"))
+	assert.Empty(t, pinnedTaskPluginChannelTypes(nil, "legacy-select"))
+}
+
+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,
+	})
+
+	assert.Empty(t, pinnedTaskPluginChannelTypes(c, "generic-select"))
+}
+
+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,
+	})
+
+	assert.Equal(t, []int{constant.ChannelTypeGemini, constant.ChannelTypeVertexAi}, pinnedTaskPluginChannelTypes(c, candidates[0].Plugin.Meta.Key))
+}
+
+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,
+	})
+
+	assert.Equal(t, []int{constant.ChannelTypeSora, constant.ChannelTypeOpenAI}, pinnedTaskPluginChannelTypes(c, "sora-select"))
+}
+
+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)
+}