mirror of
https://github.com/alibaba/higress.git
synced 2026-02-06 15:10:54 +08:00
special handling for cases where extracted content is empty and add unit test (#3251)
This commit is contained in:
rinfx
GitHub
@@ -0,0 +1,634 @@
|
||||
// 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 common
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
cfg "github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSha256Hex(t *testing.T) {
|
||||
t.Run("empty input", func(t *testing.T) {
|
||||
result := sha256Hex([]byte(""))
|
||||
// SHA256 of empty string
|
||||
expected := "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
|
||||
require.Equal(t, expected, result)
|
||||
})
|
||||
|
||||
t.Run("simple string", func(t *testing.T) {
|
||||
result := sha256Hex([]byte("hello"))
|
||||
expected := "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
|
||||
require.Equal(t, expected, result)
|
||||
})
|
||||
|
||||
t.Run("unicode string", func(t *testing.T) {
|
||||
result := sha256Hex([]byte("你好"))
|
||||
// Just verify it returns a valid hex string
|
||||
require.Len(t, result, 64)
|
||||
_, err := hex.DecodeString(result)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestHmac256(t *testing.T) {
|
||||
t.Run("valid hmac", func(t *testing.T) {
|
||||
key := []byte("test-key")
|
||||
message := "test-message"
|
||||
result, err := hmac256(key, message)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, result, 32) // SHA256 produces 32 bytes
|
||||
})
|
||||
|
||||
t.Run("empty key", func(t *testing.T) {
|
||||
key := []byte("")
|
||||
message := "test-message"
|
||||
result, err := hmac256(key, message)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, result, 32)
|
||||
})
|
||||
|
||||
t.Run("empty message", func(t *testing.T) {
|
||||
key := []byte("test-key")
|
||||
message := ""
|
||||
result, err := hmac256(key, message)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, result, 32)
|
||||
})
|
||||
|
||||
t.Run("verify hmac consistency", func(t *testing.T) {
|
||||
key := []byte("test-key")
|
||||
message := "test-message"
|
||||
result1, err1 := hmac256(key, message)
|
||||
result2, err2 := hmac256(key, message)
|
||||
require.NoError(t, err1)
|
||||
require.NoError(t, err2)
|
||||
require.Equal(t, result1, result2)
|
||||
})
|
||||
}
|
||||
|
||||
func TestPercentCode(t *testing.T) {
|
||||
t.Run("replace plus sign", func(t *testing.T) {
|
||||
input := "test+value"
|
||||
result := percentCode(input)
|
||||
require.Equal(t, "test%20value", result)
|
||||
})
|
||||
|
||||
t.Run("replace asterisk", func(t *testing.T) {
|
||||
input := "test*value"
|
||||
result := percentCode(input)
|
||||
require.Equal(t, "test%2Avalue", result)
|
||||
})
|
||||
|
||||
t.Run("replace tilde encoding", func(t *testing.T) {
|
||||
input := "test%7Evalue"
|
||||
result := percentCode(input)
|
||||
require.Equal(t, "test~value", result)
|
||||
})
|
||||
|
||||
t.Run("multiple replacements", func(t *testing.T) {
|
||||
input := "test+value*test%7E"
|
||||
result := percentCode(input)
|
||||
require.Equal(t, "test%20value%2Atest~", result)
|
||||
})
|
||||
|
||||
t.Run("no replacements needed", func(t *testing.T) {
|
||||
input := "test-value"
|
||||
result := percentCode(input)
|
||||
require.Equal(t, "test-value", result)
|
||||
})
|
||||
}
|
||||
|
||||
func TestProcessObject(t *testing.T) {
|
||||
t.Run("simple string value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
processObject(result, "key", "value")
|
||||
require.Equal(t, "value", result["key"])
|
||||
})
|
||||
|
||||
t.Run("simple int value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
processObject(result, "key", 123)
|
||||
require.Equal(t, "123", result["key"])
|
||||
})
|
||||
|
||||
t.Run("nil value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
processObject(result, "key", nil)
|
||||
require.Empty(t, result)
|
||||
})
|
||||
|
||||
t.Run("map value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := map[string]interface{}{
|
||||
"subkey1": "value1",
|
||||
"subkey2": "value2",
|
||||
}
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "value1", result["key.subkey1"])
|
||||
require.Equal(t, "value2", result["key.subkey2"])
|
||||
})
|
||||
|
||||
t.Run("array value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := []interface{}{"item1", "item2", "item3"}
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "item1", result["key.1"])
|
||||
require.Equal(t, "item2", result["key.2"])
|
||||
require.Equal(t, "item3", result["key.3"])
|
||||
})
|
||||
|
||||
t.Run("nested map", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := map[string]interface{}{
|
||||
"level1": map[string]interface{}{
|
||||
"level2": "value",
|
||||
},
|
||||
}
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "value", result["key.level1.level2"])
|
||||
})
|
||||
|
||||
t.Run("nested array", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := []interface{}{
|
||||
[]interface{}{"nested1", "nested2"},
|
||||
}
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "nested1", result["key.1.1"])
|
||||
require.Equal(t, "nested2", result["key.1.2"])
|
||||
})
|
||||
|
||||
t.Run("key with leading dot", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
processObject(result, ".key", "value")
|
||||
require.Equal(t, "value", result["key"])
|
||||
})
|
||||
|
||||
t.Run("byte array value", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := []byte("test")
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "test", result["key"])
|
||||
})
|
||||
|
||||
t.Run("complex nested structure", func(t *testing.T) {
|
||||
result := make(map[string]interface{})
|
||||
input := map[string]interface{}{
|
||||
"array": []interface{}{
|
||||
map[string]interface{}{
|
||||
"item": "value",
|
||||
},
|
||||
},
|
||||
}
|
||||
processObject(result, "key", input)
|
||||
require.Equal(t, "value", result["key.array.1.item"])
|
||||
})
|
||||
}
|
||||
|
||||
func TestFormDataToString(t *testing.T) {
|
||||
t.Run("simple map", func(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"key1": "value1",
|
||||
"key2": "value2",
|
||||
}
|
||||
result := formDataToString(input)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, *result, "key1=value1")
|
||||
require.Contains(t, *result, "key2=value2")
|
||||
})
|
||||
|
||||
t.Run("map with array", func(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"key": []interface{}{"item1", "item2"},
|
||||
}
|
||||
result := formDataToString(input)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, *result, "key.1=item1")
|
||||
require.Contains(t, *result, "key.2=item2")
|
||||
})
|
||||
|
||||
t.Run("map with nested map", func(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"key": map[string]interface{}{
|
||||
"subkey": "value",
|
||||
},
|
||||
}
|
||||
result := formDataToString(input)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, *result, "key.subkey=value")
|
||||
})
|
||||
|
||||
t.Run("empty map", func(t *testing.T) {
|
||||
input := map[string]interface{}{}
|
||||
result := formDataToString(input)
|
||||
require.NotNil(t, result)
|
||||
require.Empty(t, *result)
|
||||
})
|
||||
|
||||
t.Run("map with nil value", func(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"key1": "value1",
|
||||
"key2": nil,
|
||||
}
|
||||
result := formDataToString(input)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, *result, "key1=value1")
|
||||
require.NotContains(t, *result, "key2")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGenerateRequestForText(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "",
|
||||
}
|
||||
|
||||
t.Run("basic text request", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.True(t, strings.HasPrefix(path, "?"))
|
||||
require.Contains(t, path, "Service=llm_query_moderation")
|
||||
|
||||
require.NotEmpty(t, headers)
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Equal(t, "TextModerationPlus", headerMap["x-acs-action"])
|
||||
require.Equal(t, "2022-03-02", headerMap["x-acs-version"])
|
||||
require.Equal(t, "application/x-www-form-urlencoded", headerMap["content-type"])
|
||||
require.Equal(t, cfg.AliyunUserAgent, headerMap["User-Agent"])
|
||||
require.Contains(t, headerMap, "Authorization")
|
||||
require.Contains(t, headerMap, "x-acs-date")
|
||||
require.Contains(t, headerMap, "x-acs-signature-nonce")
|
||||
require.Contains(t, headerMap, "x-acs-content-sha256")
|
||||
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ServiceParameters")
|
||||
// Body is URL encoded, so decode it to check content
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, "test content")
|
||||
require.Contains(t, decodedBody, "test-session-id")
|
||||
require.Contains(t, decodedBody, cfg.AliyunUserAgent)
|
||||
})
|
||||
|
||||
t.Run("request with security token", func(t *testing.T) {
|
||||
configWithToken := config
|
||||
configWithToken.Token = "test-token"
|
||||
path, headers, body := GenerateRequestForText(
|
||||
configWithToken,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Equal(t, "test-token", headerMap["x-acs-security-token"])
|
||||
require.NotEmpty(t, body)
|
||||
})
|
||||
|
||||
t.Run("empty content", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ServiceParameters")
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, `"content":""`)
|
||||
})
|
||||
|
||||
t.Run("different check service", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_response_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
require.Contains(t, path, "Service=llm_response_moderation")
|
||||
require.NotEmpty(t, headers)
|
||||
require.NotEmpty(t, body)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGenerateRequestForImage(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "",
|
||||
}
|
||||
|
||||
t.Run("image request with URL", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForImage(
|
||||
config,
|
||||
"MultiModalGuard",
|
||||
"llm_image_moderation",
|
||||
"https://example.com/image.jpg",
|
||||
"",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.True(t, strings.HasPrefix(path, "?"))
|
||||
require.Contains(t, path, "Service=llm_image_moderation")
|
||||
|
||||
require.NotEmpty(t, headers)
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Equal(t, "MultiModalGuard", headerMap["x-acs-action"])
|
||||
require.Equal(t, "2022-03-02", headerMap["x-acs-version"])
|
||||
require.Equal(t, "application/x-www-form-urlencoded", headerMap["content-type"])
|
||||
require.Equal(t, cfg.AliyunUserAgent, headerMap["User-Agent"])
|
||||
require.Contains(t, headerMap, "Authorization")
|
||||
require.Contains(t, headerMap, "x-acs-date")
|
||||
require.Contains(t, headerMap, "x-acs-signature-nonce")
|
||||
require.Contains(t, headerMap, "x-acs-content-sha256")
|
||||
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ServiceParameters")
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, "https://example.com/image.jpg")
|
||||
require.Contains(t, decodedBody, cfg.AliyunUserAgent)
|
||||
})
|
||||
|
||||
t.Run("image request with base64", func(t *testing.T) {
|
||||
base64Data := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
path, headers, body := GenerateRequestForImage(
|
||||
config,
|
||||
"MultiModalGuard",
|
||||
"llm_image_moderation",
|
||||
"",
|
||||
base64Data,
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ImageBase64Str")
|
||||
// Base64 data is URL encoded, decode to check
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, base64Data)
|
||||
})
|
||||
|
||||
t.Run("image request with both URL and base64", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForImage(
|
||||
config,
|
||||
"MultiModalGuard",
|
||||
"llm_image_moderation",
|
||||
"https://example.com/image.jpg",
|
||||
"base64data",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ImageBase64Str")
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, "https://example.com/image.jpg")
|
||||
require.Contains(t, decodedBody, "base64data")
|
||||
})
|
||||
|
||||
t.Run("image request with security token", func(t *testing.T) {
|
||||
configWithToken := config
|
||||
configWithToken.Token = "test-token"
|
||||
path, headers, body := GenerateRequestForImage(
|
||||
configWithToken,
|
||||
"MultiModalGuard",
|
||||
"llm_image_moderation",
|
||||
"https://example.com/image.jpg",
|
||||
"",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Equal(t, "test-token", headerMap["x-acs-security-token"])
|
||||
require.NotEmpty(t, body)
|
||||
})
|
||||
|
||||
t.Run("empty image URL and base64", func(t *testing.T) {
|
||||
path, headers, body := GenerateRequestForImage(
|
||||
config,
|
||||
"MultiModalGuard",
|
||||
"llm_image_moderation",
|
||||
"",
|
||||
"",
|
||||
)
|
||||
|
||||
require.NotEmpty(t, path)
|
||||
require.NotEmpty(t, headers)
|
||||
require.NotEmpty(t, body)
|
||||
bodyStr := string(body)
|
||||
require.Contains(t, bodyStr, "ServiceParameters")
|
||||
decodedBody, err := url.QueryUnescape(bodyStr)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, decodedBody, cfg.AliyunUserAgent)
|
||||
require.NotContains(t, decodedBody, "imageUrls")
|
||||
require.NotContains(t, decodedBody, "ImageBase64Str")
|
||||
})
|
||||
}
|
||||
|
||||
func TestNewRequest(t *testing.T) {
|
||||
// Test newRequest indirectly through GenerateRequestForText
|
||||
// Since it's a private function, we test it through public API
|
||||
t.Run("request structure", func(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "",
|
||||
}
|
||||
|
||||
path, headers, _ := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test",
|
||||
"session-id",
|
||||
)
|
||||
|
||||
// Verify that newRequest was called correctly by checking headers
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
// Verify headers set by newRequest
|
||||
require.Equal(t, "TextModerationPlus", headerMap["x-acs-action"])
|
||||
require.Equal(t, "2022-03-02", headerMap["x-acs-version"])
|
||||
require.Contains(t, headerMap, "x-acs-date")
|
||||
require.Contains(t, headerMap, "x-acs-signature-nonce")
|
||||
require.NotEmpty(t, path)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetAuthorization(t *testing.T) {
|
||||
// Test getAuthorization indirectly through GenerateRequestForText
|
||||
// Since it's a private function, we test it through public API
|
||||
t.Run("authorization header format", func(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "",
|
||||
}
|
||||
|
||||
_, headers, _ := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
authHeader := headerMap["Authorization"]
|
||||
require.NotEmpty(t, authHeader)
|
||||
require.Contains(t, authHeader, "ACS3-HMAC-SHA256")
|
||||
require.Contains(t, authHeader, "Credential=test-ak")
|
||||
require.Contains(t, authHeader, "SignedHeaders=")
|
||||
require.Contains(t, authHeader, "Signature=")
|
||||
|
||||
// Verify content SHA256 is set
|
||||
require.Contains(t, headerMap, "x-acs-content-sha256")
|
||||
require.Len(t, headerMap["x-acs-content-sha256"], 64) // SHA256 hex string length
|
||||
})
|
||||
|
||||
t.Run("authorization with security token", func(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "test-token",
|
||||
}
|
||||
|
||||
_, headers, _ := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
headerMap := make(map[string]string)
|
||||
for _, h := range headers {
|
||||
headerMap[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Equal(t, "test-token", headerMap["x-acs-security-token"])
|
||||
require.Contains(t, headerMap, "Authorization")
|
||||
})
|
||||
|
||||
t.Run("authorization signature consistency", func(t *testing.T) {
|
||||
config := cfg.AISecurityConfig{
|
||||
Host: "security.example.com",
|
||||
AK: "test-ak",
|
||||
SK: "test-sk",
|
||||
Token: "",
|
||||
}
|
||||
|
||||
// Generate two requests with same content
|
||||
_, headers1, body1 := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
_, headers2, body2 := GenerateRequestForText(
|
||||
config,
|
||||
"TextModerationPlus",
|
||||
"llm_query_moderation",
|
||||
"test content",
|
||||
"test-session-id",
|
||||
)
|
||||
|
||||
// Bodies should be the same (except for sessionId which is random)
|
||||
require.NotEmpty(t, body1)
|
||||
require.NotEmpty(t, body2)
|
||||
|
||||
// Headers should have authorization
|
||||
headerMap1 := make(map[string]string)
|
||||
for _, h := range headers1 {
|
||||
headerMap1[h[0]] = h[1]
|
||||
}
|
||||
|
||||
headerMap2 := make(map[string]string)
|
||||
for _, h := range headers2 {
|
||||
headerMap2[h[0]] = h[1]
|
||||
}
|
||||
|
||||
require.Contains(t, headerMap1, "Authorization")
|
||||
require.Contains(t, headerMap2, "Authorization")
|
||||
// Signatures will be different due to nonce and timestamp, but format should be same
|
||||
require.Contains(t, headerMap1["Authorization"], "ACS3-HMAC-SHA256")
|
||||
require.Contains(t, headerMap2["Authorization"], "ACS3-HMAC-SHA256")
|
||||
})
|
||||
}
|
||||
@@ -100,10 +100,11 @@ func HandleTextGenerationStreamingResponseBody(ctx wrapper.HttpContext, config c
|
||||
break
|
||||
}
|
||||
}
|
||||
// if streaming body has reasoning_content, buffer maybe empty
|
||||
// case 1: streaming body has reasoning_content, part of buffer maybe empty
|
||||
// case 2: streaming body has toolcall result, part of buffer maybe empty
|
||||
log.Debugf("current content piece: %s", buffer)
|
||||
if len(buffer) == 0 {
|
||||
return
|
||||
buffer = "[empty content]"
|
||||
}
|
||||
ctx.SetContext("during_call", true)
|
||||
log.Debugf("current content piece: %s", buffer)
|
||||
|
||||
@@ -384,6 +384,76 @@ func TestOnHttpResponseHeaders(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestOnHttpResponseBody(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// 测试响应体安全检查通过
|
||||
t.Run("response body security check pass", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 先设置请求头
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
// 设置响应头
|
||||
host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
// 设置响应体
|
||||
body := `{"choices": [{"message": {"role": "assistant", "content": "Hello, how can I help you?"}}]}`
|
||||
action := host.CallOnHttpResponseBody([]byte(body))
|
||||
|
||||
// 应该返回ActionPause，等待安全检查结果
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
// 模拟安全检查服务响应（通过）
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-123", "Data": {"RiskLevel": "low"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
action = host.GetHttpStreamAction()
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试空响应内容
|
||||
t.Run("empty response content", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 先设置请求头
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
// 设置响应头
|
||||
host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
// 设置空内容的响应体
|
||||
body := `{"choices": [{"message": {"role": "assistant", "content": ""}}]}`
|
||||
action := host.CallOnHttpResponseBody([]byte(body))
|
||||
|
||||
// 空内容应该直接通过
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestRiskLevelFunctions(t *testing.T) {
|
||||
// 测试风险等级转换函数
|
||||
t.Run("risk level conversion", func(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user