new-api/controller/kling_aiping_native_test.go

package controller

import (
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/model"
	klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func TestKlingAipingSubmitPreparationUsesModelNameBeforeModel(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{
		"model":"Kling-V1.6",
		"model_name":"Kling-V2.6",
		"prompt":"prompt"
	}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
	modelName := resolveKlingAipingModel(payload, route)
	relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{
		Model:    modelName,
		Prompt:   stringFromMap(payload, "prompt"),
		Metadata: payload,
	})

	stored, err := relaycommon.GetTaskRequest(c)
	require.NoError(t, err)
	require.Equal(t, "Kling-V2.6", stored.Model)
	require.Equal(t, "prompt", stored.Prompt)
	require.Equal(t, "Kling-V1.6", stored.Metadata["model"])
	require.Equal(t, "Kling-V2.6", stored.Metadata["model_name"])
}

func TestKlingAipingSubmitPreparationDoesNotLockChannel(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/general/custom-voices", `{
		"voice_url":"https://example.com/voice.mp3",
		"voice_name":"voice"
	}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/general/custom-voices", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
	info.OriginModelName = resolveKlingAipingModel(payload, route)
	info.Action = route.Action
	relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{
		Model:    info.OriginModelName,
		Metadata: payload,
	})

	require.Equal(t, klingaiping.ModelCustomVoices, info.OriginModelName)
	require.Equal(t, klingaiping.ActionVoicesCreate, info.Action)
	require.Nil(t, info.LockedChannel)
}

func TestConfigureKlingAipingTaskRelayInfoForcesChannelSelection(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt"}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	configureKlingAipingTaskRelayInfo(c, info, route, payload)

	require.NotNil(t, info.ChannelMeta)
	require.Zero(t, info.ChannelMeta.ChannelType)
	require.Equal(t, "Kling-V2.6", info.OriginModelName)
	require.Equal(t, klingaiping.ActionText2Video, info.Action)
	require.Nil(t, info.LockedChannel)
}

func TestConfigureKlingAipingTaskRelayInfoStoresDuration(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt","duration":5}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	configureKlingAipingTaskRelayInfo(c, info, route, payload)

	stored, err := relaycommon.GetTaskRequest(c)
	require.NoError(t, err)
	require.Equal(t, 5, stored.Duration)
	require.Empty(t, stored.Seconds)
}

func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`)
	c.Request.URL.Path = "/v1/videos/text2video"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.ElementsMatch(t,
		service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling),
		allowedTaskChannelTypesForRequest(c),
	)

	c = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`)
	c.Request.URL.Path = "/v1/videos/video-extend"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.ElementsMatch(t,
		service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling),
		allowedTaskChannelTypesForRequest(c),
	)

	c = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`)
	c.Request.URL.Path = "/v1/videos/video_123/remix"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.Nil(t, allowedTaskChannelTypesForRequest(c))
}

func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) {
	task := &model.Task{
		TaskID:    "task_public",
		Status:    model.TaskStatusSuccess,
		CreatedAt: 100,
		UpdatedAt: 200,
		Data: []byte(`{
			"code":0,
			"aiping_id":"internal",
			"data":{
				"task_id":"899333358055493641",
				"task_status":"succeed",
				"task_result":{"videos":[{"id":"v1","url":"https://example.com/video.mp4","duration":"5.041"}]}
			}
		}`),
	}

	data := taskDataObject(task)
	require.Equal(t, "task_public", data["task_id"])
	taskResult := data["task_result"].(map[string]any)
	videos := taskResult["videos"].([]any)
	require.Equal(t, "", videos[0].(map[string]any)["watermark_url"])
}

func TestDoKlingAipingProxyRequestUsesSelectedContextKey(t *testing.T) {
	service.InitHttpClient()
	var gotAuth string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"code":0,"message":"success","data":[]}`))
	}))
	defer server.Close()

	c := newControllerJSONContext(t, "/v1/general/advanced-presets-elements", ``)
	common.SetContextKey(c, constant.ContextKeyChannelKey, "selected-key")
	route, ok := klingaiping.FindRoute(http.MethodGet, "/v1/general/advanced-presets-elements", klingaiping.RouteKindProxy)
	require.True(t, ok)
	channel := &model.Channel{
		Key:     "raw-channel-key",
		BaseURL: common.GetPointer(server.URL),
	}

	resp, err := doKlingAipingProxyRequest(c, route, channel)
	require.NoError(t, err)
	defer resp.Body.Close()

	require.Equal(t, "Bearer selected-key", gotAuth)
}

func TestKlingAipingNativeProxyPersistsFallbackBinding(t *testing.T) {
	service.InitHttpClient()
	db := setupKlingAipingNativeProxyDB(t)
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"code":0,"message":"success","data":[]}`))
	}))
	defer upstream.Close()
	createKlingAipingProxyChannelForTest(t, db, 59, "default", "proxy-key", upstream.URL)
	createKlingAipingProxyAbilityForTest(t, db, "default", klingaiping.ModelKlingAdvancedElements, 59, true)

	w := httptest.NewRecorder()
	_, engine := gin.CreateTestContext(w)
	engine.GET("/v1/general/advanced-presets-elements", func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
		KlingAipingNativeProxy(c)
	})

	engine.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/general/advanced-presets-elements", nil))

	require.Equal(t, http.StatusOK, w.Code)
	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default")
	require.NoError(t, err)
	require.NotNil(t, binding)
	require.Equal(t, 59, binding.ChannelId)
}

func TestCopyProxyResponseNormalizesMsgError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusUnauthorized,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(`{"code":401,"msg":"unauthorized","data":null}`)),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusUnauthorized, w.Code)
	require.Contains(t, w.Body.String(), `"message":"unauthorized"`)
	require.NotContains(t, w.Body.String(), `"msg"`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestCopyProxyResponseNormalizesPlainTextError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusMethodNotAllowed,
		Header:     http.Header{"Content-Type": []string{"text/plain"}},
		Body:       io.NopCloser(strings.NewReader("Method Not Allowed")),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusMethodNotAllowed, w.Code)
	require.Contains(t, w.Body.String(), `"message":"Method Not Allowed"`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestCopyProxyResponseNormalizesDetailMessageError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusServiceUnavailable,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(`{"detail":{"message":"not found","error_type":"not_found"},"aiping_id":"internal"}`)),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusServiceUnavailable, w.Code)
	require.Contains(t, w.Body.String(), `"message":"not found"`)
	require.NotContains(t, w.Body.String(), `map[`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestNormalizeKlingAipingTaskErrorMessageExtractsUpstreamJSONMessage(t *testing.T) {
	taskErr := &dto.TaskError{
		Code:       "fail_to_fetch_task",
		Message:    `{"code":400,"message":"ERROR: image download failed","request_id":"upstream"}`,
		StatusCode: http.StatusBadRequest,
	}

	normalizeKlingAipingTaskError(taskErr)

	require.Equal(t, "ERROR: image download failed", taskErr.Message)
}

func TestParseKlingAipingPageBoundaries(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1001&pageSize=30", `{}`)
	c.Request.URL.RawQuery = "pageNum=1001&pageSize=30"
	_, _, err := parseKlingAipingPage(c)
	require.ErrorContains(t, err, "pageNum")

	c = newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1&pageSize=501", `{}`)
	c.Request.URL.RawQuery = "pageNum=1&pageSize=501"
	_, _, err = parseKlingAipingPage(c)
	require.ErrorContains(t, err, "pageSize")
}

func setupKlingAipingNativeProxyDB(t *testing.T) *gorm.DB {
	t.Helper()

	oldDB := model.DB
	oldSQLitePath := common.SQLitePath
	oldMemoryCacheEnabled := common.MemoryCacheEnabled
	oldIsMasterNode := common.IsMasterNode
	oldUsingSQLite := common.UsingSQLite
	oldUsingMySQL := common.UsingMySQL
	oldUsingPostgreSQL := common.UsingPostgreSQL
	oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")

	common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
	common.MemoryCacheEnabled = false
	common.IsMasterNode = false
	common.UsingSQLite = false
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	require.NoError(t, os.Setenv("SQL_DSN", "local"))
	require.NoError(t, model.InitDB())

	model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
	db := model.DB
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}))

	t.Cleanup(func() {
		_ = sqlDB.Close()
		model.DB = oldDB
		common.SQLitePath = oldSQLitePath
		common.MemoryCacheEnabled = oldMemoryCacheEnabled
		common.IsMasterNode = oldIsMasterNode
		common.UsingSQLite = oldUsingSQLite
		common.UsingMySQL = oldUsingMySQL
		common.UsingPostgreSQL = oldUsingPostgreSQL
		if hadSQLDSN {
			_ = os.Setenv("SQL_DSN", oldSQLDSN)
		} else {
			_ = os.Unsetenv("SQL_DSN")
		}
	})

	return db
}

func createKlingAipingProxyChannelForTest(t *testing.T, db *gorm.DB, id int, group string, key string, baseURL string) {
	t.Helper()

	priority := int64(id)
	weight := uint(10)
	autoBan := 1
	require.NoError(t, db.Create(&model.Channel{
		Id:       id,
		Type:     constant.ChannelTypeKlingAiping,
		Key:      key,
		Status:   common.ChannelStatusEnabled,
		Name:     "kling-aiping-proxy",
		Group:    group,
		Models:   klingaiping.ModelKlingAdvancedElements,
		BaseURL:  common.GetPointer(baseURL),
		Priority: &priority,
		Weight:   &weight,
		AutoBan:  &autoBan,
	}).Error)
}

func createKlingAipingProxyAbilityForTest(t *testing.T, db *gorm.DB, group string, modelName string, channelId int, enabled bool) {
	t.Helper()

	priority := int64(channelId)
	require.NoError(t, db.Create(&model.Ability{
		Group:     group,
		Model:     modelName,
		ChannelId: channelId,
		Enabled:   enabled,
		Priority:  &priority,
		Weight:    10,
	}).Error)
}