mirror of
https://github.com/alibaba/higress.git
synced 2026-06-08 12:17:28 +08:00
feat(wasm-go): add wasm go plugin unit test and ci workflow (#2809)
This commit is contained in:
Jingze
GitHub
390
plugins/wasm-go/extensions/chatgpt-proxy/main_test.go
Normal file
390
plugins/wasm-go/extensions/chatgpt-proxy/main_test.go
Normal file
@@ -0,0 +1,390 @@
|
||||
// Copyright (c) 2024 Alibaba Group Holding Ltd.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/higress-group/wasm-go/pkg/test"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 测试配置：基本配置
|
||||
var basicConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "prompt",
|
||||
"model": "text-davinci-003",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：自定义模型配置
|
||||
var customModelConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "text",
|
||||
"model": "curie",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：自定义提示参数配置
|
||||
var customPromptParamConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "question",
|
||||
"model": "text-davinci-003",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：自定义 ChatGPT URI 配置
|
||||
var customUriConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "prompt",
|
||||
"model": "text-davinci-003",
|
||||
"chatgptUri": "https://custom-ai.example.com/v1/chat/completions",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：自定义 Human ID 和 AI ID 配置
|
||||
var customIdsConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "prompt",
|
||||
"model": "text-davinci-003",
|
||||
"HumainId": "User:",
|
||||
"AIId": "Assistant:",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：无效配置（缺少 API Key）
|
||||
var invalidConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"promptParam": "prompt",
|
||||
"model": "text-davinci-003",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置：无效 URI 配置
|
||||
var invalidUriConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"apiKey": "sk-test123456789",
|
||||
"promptParam": "prompt",
|
||||
"model": "text-davinci-003",
|
||||
"chatgptUri": "://invalid-uri",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
func TestParseConfig(t *testing.T) {
|
||||
test.RunGoTest(t, func(t *testing.T) {
|
||||
// 测试基本配置解析
|
||||
t.Run("basic config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
// 测试自定义模型配置解析
|
||||
t.Run("custom model config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(customModelConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
// 测试自定义提示参数配置解析
|
||||
t.Run("custom prompt param config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(customPromptParamConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
// 测试自定义 URI 配置解析
|
||||
t.Run("custom uri config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(customUriConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
// 测试自定义 ID 配置解析
|
||||
t.Run("custom ids config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(customIdsConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
// 测试无效配置（缺少 API Key）
|
||||
t.Run("invalid config - missing api key", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(invalidConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusFailed, status)
|
||||
})
|
||||
|
||||
// 测试无效 URI 配置
|
||||
t.Run("invalid config - invalid uri", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(invalidUriConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusFailed, status)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// 测试基本请求头处理（带查询参数）
|
||||
t.Run("basic request headers with query params", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头，包含查询参数
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?prompt=Hello, how are you?"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟外部 AI 服务响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "200"},
|
||||
}, []byte(`{"choices":[{"text":"I'm doing well, thank you for asking!"}]}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(200), response.StatusCode)
|
||||
require.Equal(t, `{"choices":[{"text":"I'm doing well, thank you for asking!"}]}`, string(response.Data))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试自定义提示参数请求头处理
|
||||
t.Run("custom prompt param request headers", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(customPromptParamConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头，使用自定义提示参数
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?question=What is the weather like?"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟外部 AI 服务响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "200"},
|
||||
}, []byte(`{"choices":[{"text":"I don't have access to real-time weather information."}]}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(200), response.StatusCode)
|
||||
require.Equal(t, `{"choices":[{"text":"I don't have access to real-time weather information."}]}`, string(response.Data))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试缺少查询参数的情况
|
||||
t.Run("missing query params", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头，不包含查询参数
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 应该返回 ActionContinue，因为缺少查询参数
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试缺少提示参数的情况
|
||||
t.Run("missing prompt param", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头，包含查询参数但不包含提示参数
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?other=value"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 应该返回 ActionContinue，因为缺少提示参数
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试空提示参数的情况
|
||||
t.Run("empty prompt param", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头，包含空的提示参数
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?prompt="},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟外部 AI 服务响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "200"},
|
||||
}, []byte(`{"choices":[{"text":"Empty prompt response"}]}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(200), response.StatusCode)
|
||||
require.Equal(t, `{"choices":[{"text":"Empty prompt response"}]}`, string(response.Data))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试外部服务调用成功的情况
|
||||
t.Run("external service call success", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?prompt=Tell me a joke"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟外部 AI 服务成功响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "200"},
|
||||
}, []byte(`{"choices":[{"text":"Why don't scientists trust atoms? Because they make up everything!"}]}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(200), response.StatusCode)
|
||||
require.Equal(t, `{"choices":[{"text":"Why don't scientists trust atoms? Because they make up everything!"}]}`, string(response.Data))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试外部服务调用失败的情况
|
||||
t.Run("external service call failure", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 设置请求头
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?prompt=Hello"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟外部 AI 服务失败响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "429"},
|
||||
}, []byte(`{"error":"Rate limit exceeded"}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(429), response.StatusCode)
|
||||
require.Equal(t, `{"error":"Rate limit exceeded"}`, string(response.Data))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestCompleteFlow(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
t.Run("complete chatgpt proxy flow", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 1. 处理请求头
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/chat?prompt=What is artificial intelligence?"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
|
||||
// 由于需要调用外部 AI 服务，应该返回 ActionPause
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 2. 模拟外部 AI 服务响应
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{":status", "200"},
|
||||
}, []byte(`{"choices":[{"text":"Artificial Intelligence (AI) is a branch of computer science that aims to create systems capable of performing tasks that typically require human intelligence."}]}`))
|
||||
|
||||
response := host.GetLocalResponse()
|
||||
require.Equal(t, uint32(200), response.StatusCode)
|
||||
require.Equal(t, `{"choices":[{"text":"Artificial Intelligence (AI) is a branch of computer science that aims to create systems capable of performing tasks that typically require human intelligence."}]}`, string(response.Data))
|
||||
|
||||
// 3. 完成请求
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user