mirror of
https://github.com/alibaba/higress.git
synced 2026-05-21 19:27:28 +08:00
419 lines
15 KiB
Go
419 lines
15 KiB
Go
package test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
|
"github.com/higress-group/wasm-go/pkg/test"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
var klingOfficialConfig = func() json.RawMessage {
|
|
data, _ := json.Marshal(map[string]interface{}{
|
|
"provider": map[string]interface{}{
|
|
"type": "kling",
|
|
"klingAccessKey": "kling-ak-test",
|
|
"klingSecretKey": "kling-sk-test",
|
|
"klingTokenRefreshAhead": 60,
|
|
"modelMapping": map[string]string{
|
|
"client-video": "kling-v2-1",
|
|
},
|
|
},
|
|
})
|
|
return data
|
|
}()
|
|
|
|
var klingGatewayConfig = func() json.RawMessage {
|
|
data, _ := json.Marshal(map[string]interface{}{
|
|
"provider": map[string]interface{}{
|
|
"type": "kling",
|
|
"apiTokens": []string{"gateway-token"},
|
|
"providerDomain": "api.302.ai",
|
|
"providerBasePath": "/klingai",
|
|
"modelMapping": map[string]string{
|
|
"client-video": "kling-v2-1",
|
|
},
|
|
},
|
|
})
|
|
return data
|
|
}()
|
|
|
|
var klingGatewayCustomImageRetrieveConfig = func() json.RawMessage {
|
|
data, _ := json.Marshal(map[string]interface{}{
|
|
"provider": map[string]interface{}{
|
|
"type": "kling",
|
|
"apiTokens": []string{"gateway-token"},
|
|
"providerDomain": "api.302.ai",
|
|
"providerBasePath": "/klingai",
|
|
"modelMapping": map[string]string{
|
|
"client-video": "kling-v2-1",
|
|
},
|
|
"capabilities": map[string]string{
|
|
"openai/v1/videos": "/gateway/text2video?mode=text",
|
|
"kling/v1/image2video": "/gateway/image2video?mode=image",
|
|
"kling/v1/retrieveimagevideo": "/gateway/image-tasks/{video_id}?version=1",
|
|
},
|
|
},
|
|
})
|
|
return data
|
|
}()
|
|
|
|
var klingOriginalConfig = func() json.RawMessage {
|
|
data, _ := json.Marshal(map[string]interface{}{
|
|
"provider": map[string]interface{}{
|
|
"type": "kling",
|
|
"apiTokens": []string{"gateway-token"},
|
|
"providerDomain": "api.302.ai",
|
|
"providerBasePath": "/klingai",
|
|
"protocol": "original",
|
|
},
|
|
})
|
|
return data
|
|
}()
|
|
|
|
func RunKlingParseConfigTests(t *testing.T) {
|
|
test.RunGoTest(t, func(t *testing.T) {
|
|
t.Run("kling official config", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingOfficialConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
config, err := host.GetMatchConfig()
|
|
require.NoError(t, err)
|
|
require.NotNil(t, config)
|
|
})
|
|
|
|
t.Run("kling gateway config", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
config, err := host.GetMatchConfig()
|
|
require.NoError(t, err)
|
|
require.NotNil(t, config)
|
|
})
|
|
})
|
|
}
|
|
|
|
func RunKlingOnHttpRequestHeadersTests(t *testing.T) {
|
|
test.RunTest(t, func(t *testing.T) {
|
|
t.Run("official mode sets jwt bearer and default host", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingOfficialConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
require.Equal(t, types.HeaderStopIteration, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority")
|
|
require.True(t, hasHost)
|
|
require.Equal(t, "api-singapore.klingai.com", hostValue)
|
|
|
|
authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization")
|
|
require.True(t, hasAuth)
|
|
require.True(t, strings.HasPrefix(authValue, "Bearer "))
|
|
require.Len(t, strings.Split(strings.TrimPrefix(authValue, "Bearer "), "."), 3)
|
|
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/v1/videos/text2video", pathValue)
|
|
})
|
|
|
|
t.Run("providerDomain and providerBasePath apply to gateway mode", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
require.Equal(t, types.HeaderStopIteration, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority")
|
|
require.True(t, hasHost)
|
|
require.Equal(t, "api.302.ai", hostValue)
|
|
|
|
authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization")
|
|
require.True(t, hasAuth)
|
|
require.Equal(t, "Bearer gateway-token", authValue)
|
|
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/text2video", pathValue)
|
|
})
|
|
|
|
t.Run("retrieve video query path is mapped under providerBasePath", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/task-123?with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/text2video/task-123?with_status=true", pathValue)
|
|
})
|
|
|
|
t.Run("prefixed image task query path is mapped to image endpoint", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/kling-i2v-task-123?with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/image2video/task-123?with_status=true", pathValue)
|
|
})
|
|
|
|
t.Run("prefixed image task query strips task type hint", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/kling-i2v-task-123?kling_task_type=image2video&with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/image2video/task-123?with_status=true", pathValue)
|
|
})
|
|
|
|
t.Run("raw image task query path uses explicit task type hint", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/raw-task-123?kling_task_type=image2video&with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/image2video/raw-task-123?with_status=true", pathValue)
|
|
})
|
|
|
|
t.Run("raw retrieve strips unknown task type hint before fallback mapping", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/raw-task-123?kling_task_type=bad&with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/text2video/raw-task-123?with_status=true", pathValue)
|
|
})
|
|
|
|
t.Run("image retrieve path uses configured image capability", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayCustomImageRetrieveConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/kling-i2v-task-123?with_status=true"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/gateway/image-tasks/task-123?version=1&with_status=true", pathValue)
|
|
})
|
|
})
|
|
}
|
|
|
|
func RunKlingOnHttpRequestBodyTests(t *testing.T) {
|
|
test.RunTest(t, func(t *testing.T) {
|
|
t.Run("text to video keeps text endpoint and maps model_name", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos?gateway_param=1"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
|
|
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"sunrise","duration":"5"}`))
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/text2video?gateway_param=1", pathValue)
|
|
|
|
processedBody := host.GetRequestBody()
|
|
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
|
require.False(t, gjson.GetBytes(processedBody, "model").Exists())
|
|
require.Equal(t, "sunrise", gjson.GetBytes(processedBody, "prompt").String())
|
|
})
|
|
|
|
t.Run("image to video switches endpoint after body inspection", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos?gateway_param=1"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
|
|
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/image2video?gateway_param=1", pathValue)
|
|
|
|
processedBody := host.GetRequestBody()
|
|
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
|
require.Equal(t, "https://example.com/a.png", gjson.GetBytes(processedBody, "image").String())
|
|
})
|
|
|
|
t.Run("image to video uses configured image capability and merges query", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayCustomImageRetrieveConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos?gateway_param=1"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
|
|
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/gateway/image2video?mode=image&gateway_param=1", pathValue)
|
|
|
|
processedBody := host.GetRequestBody()
|
|
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
|
})
|
|
|
|
t.Run("original protocol does not expose request body handler", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingOriginalConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/image2video"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
{"Content-Length", "64"},
|
|
})
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/image2video", pathValue)
|
|
contentLengthValue, hasContentLength := test.GetHeaderValue(requestHeaders, "Content-Length")
|
|
require.True(t, hasContentLength)
|
|
require.Equal(t, "64", contentLengthValue)
|
|
})
|
|
|
|
t.Run("original protocol recognizes native retrieve video path", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingOriginalConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
action := host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos/text2video/task-123"},
|
|
{":method", "GET"},
|
|
})
|
|
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
|
|
|
requestHeaders := host.GetRequestHeaders()
|
|
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
|
require.True(t, hasPath)
|
|
require.Equal(t, "/klingai/v1/videos/text2video/task-123", pathValue)
|
|
})
|
|
})
|
|
}
|
|
|
|
func RunKlingOnHttpResponseBodyTests(t *testing.T) {
|
|
test.RunTest(t, func(t *testing.T) {
|
|
t.Run("image creation response prefixes task id", func(t *testing.T) {
|
|
host, status := test.NewTestHost(klingGatewayConfig)
|
|
defer host.Reset()
|
|
require.Equal(t, types.OnPluginStartStatusOK, status)
|
|
|
|
host.CallOnHttpRequestHeaders([][2]string{
|
|
{":authority", "example.com"},
|
|
{":path", "/v1/videos"},
|
|
{":method", "POST"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
require.NoError(t, host.SetProperty([]string{"response", "code_details"}, []byte("via_upstream")))
|
|
action = host.CallOnHttpResponseHeaders([][2]string{
|
|
{":status", "200"},
|
|
{"Content-Type", "application/json"},
|
|
})
|
|
require.Equal(t, types.ActionContinue, action)
|
|
action = host.CallOnHttpResponseBody([]byte(`{"id":"root-task","data":{"task_id":"task-123"}}`))
|
|
require.Equal(t, types.ActionContinue, action)
|
|
|
|
processedBody := host.GetResponseBody()
|
|
require.Equal(t, "root-task", gjson.GetBytes(processedBody, "id").String())
|
|
require.Equal(t, "kling-i2v-task-123", gjson.GetBytes(processedBody, "data.task_id").String())
|
|
})
|
|
})
|
|
}
|