package handler import ( "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) { tests := []struct { name string errType string message string }{ { name: "包含双引号的消息", errType: "server_error", message: `upstream returned "invalid" response`, }, { name: "包含反斜杠的消息", errType: "server_error", message: `path C:\Users\test\file.txt not found`, }, { name: "包含双引号和反斜杠的消息", errType: "upstream_error", message: `error parsing "key\value": unexpected token`, }, { name: "包含换行符的消息", errType: "server_error", message: "line1\nline2\ttab", }, { name: "普通消息", errType: "upstream_error", message: "Upstream service temporarily unavailable", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(http.MethodGet, "/", nil) h := &OpenAIGatewayHandler{} h.handleStreamingAwareError(c, http.StatusBadGateway, tt.errType, tt.message, true) body := w.Body.String() // 验证 SSE 格式:event: error\ndata: {JSON}\n\n assert.True(t, strings.HasPrefix(body, "event: error\n"), "应以 'event: error\\n' 开头") assert.True(t, strings.HasSuffix(body, "\n\n"), "应以 '\\n\\n' 结尾") // 提取 data 部分 lines := strings.Split(strings.TrimSuffix(body, "\n\n"), "\n") require.Len(t, lines, 2, "应有 event 行和 data 行") dataLine := lines[1] require.True(t, strings.HasPrefix(dataLine, "data: "), "第二行应以 'data: ' 开头") jsonStr := strings.TrimPrefix(dataLine, "data: ") // 验证 JSON 合法性 var parsed map[string]any err := json.Unmarshal([]byte(jsonStr), &parsed) require.NoError(t, err, "JSON 应能被成功解析,原始 JSON: %s", jsonStr) // 验证结构 errorObj, ok := parsed["error"].(map[string]any) require.True(t, ok, "应包含 error 对象") assert.Equal(t, tt.errType, errorObj["type"]) assert.Equal(t, tt.message, errorObj["message"]) }) } } func TestOpenAIHandleStreamingAwareError_NonStreaming(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(http.MethodGet, "/", nil) h := &OpenAIGatewayHandler{} h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", "test error", false) // 非流式应返回 JSON 响应 assert.Equal(t, http.StatusBadGateway, w.Code) var parsed map[string]any err := json.Unmarshal(w.Body.Bytes(), &parsed) require.NoError(t, err) errorObj, ok := parsed["error"].(map[string]any) require.True(t, ok) assert.Equal(t, "upstream_error", errorObj["type"]) assert.Equal(t, "test error", errorObj["message"]) } func TestOpenAIEnsureForwardErrorResponse_WritesFallbackWhenNotWritten(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(http.MethodGet, "/", nil) h := &OpenAIGatewayHandler{} wrote := h.ensureForwardErrorResponse(c, false) require.True(t, wrote) require.Equal(t, http.StatusBadGateway, w.Code) var parsed map[string]any err := json.Unmarshal(w.Body.Bytes(), &parsed) require.NoError(t, err) errorObj, ok := parsed["error"].(map[string]any) require.True(t, ok) assert.Equal(t, "upstream_error", errorObj["type"]) assert.Equal(t, "Upstream request failed", errorObj["message"]) } func TestOpenAIEnsureForwardErrorResponse_DoesNotOverrideWrittenResponse(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(http.MethodGet, "/", nil) c.String(http.StatusTeapot, "already written") h := &OpenAIGatewayHandler{} wrote := h.ensureForwardErrorResponse(c, false) require.False(t, wrote) require.Equal(t, http.StatusTeapot, w.Code) assert.Equal(t, "already written", w.Body.String()) } // TestOpenAIHandler_GjsonExtraction 验证 gjson 从请求体中提取 model/stream 的正确性 func TestOpenAIHandler_GjsonExtraction(t *testing.T) { tests := []struct { name string body string wantModel string wantStream bool }{ {"正常提取", `{"model":"gpt-4","stream":true,"input":"hello"}`, "gpt-4", true}, {"stream false", `{"model":"gpt-4","stream":false}`, "gpt-4", false}, {"无 stream 字段", `{"model":"gpt-4"}`, "gpt-4", false}, {"model 缺失", `{"stream":true}`, "", true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { body := []byte(tt.body) modelResult := gjson.GetBytes(body, "model") model := "" if modelResult.Type == gjson.String { model = modelResult.String() } stream := gjson.GetBytes(body, "stream").Bool() require.Equal(t, tt.wantModel, model) require.Equal(t, tt.wantStream, stream) }) } } // TestOpenAIHandler_GjsonValidation 验证修复后的 JSON 合法性和类型校验 func TestOpenAIHandler_GjsonValidation(t *testing.T) { // 非法 JSON 被 gjson.ValidBytes 拦截 require.False(t, gjson.ValidBytes([]byte(`{invalid json`))) // model 为数字 → 类型不是 gjson.String,应被拒绝 body := []byte(`{"model":123}`) modelResult := gjson.GetBytes(body, "model") require.True(t, modelResult.Exists()) require.NotEqual(t, gjson.String, modelResult.Type) // model 为 null → 类型不是 gjson.String,应被拒绝 body2 := []byte(`{"model":null}`) modelResult2 := gjson.GetBytes(body2, "model") require.True(t, modelResult2.Exists()) require.NotEqual(t, gjson.String, modelResult2.Type) // stream 为 string → 类型既不是 True 也不是 False,应被拒绝 body3 := []byte(`{"model":"gpt-4","stream":"true"}`) streamResult := gjson.GetBytes(body3, "stream") require.True(t, streamResult.Exists()) require.NotEqual(t, gjson.True, streamResult.Type) require.NotEqual(t, gjson.False, streamResult.Type) // stream 为 int → 同上 body4 := []byte(`{"model":"gpt-4","stream":1}`) streamResult2 := gjson.GetBytes(body4, "stream") require.True(t, streamResult2.Exists()) require.NotEqual(t, gjson.True, streamResult2.Type) require.NotEqual(t, gjson.False, streamResult2.Type) } // TestOpenAIHandler_InstructionsInjection 验证 instructions 的 gjson/sjson 注入逻辑 func TestOpenAIHandler_InstructionsInjection(t *testing.T) { // 测试 1:无 instructions → 注入 body := []byte(`{"model":"gpt-4"}`) existing := gjson.GetBytes(body, "instructions").String() require.Empty(t, existing) newBody, err := sjson.SetBytes(body, "instructions", "test instruction") require.NoError(t, err) require.Equal(t, "test instruction", gjson.GetBytes(newBody, "instructions").String()) // 测试 2:已有 instructions → 不覆盖 body2 := []byte(`{"model":"gpt-4","instructions":"existing"}`) existing2 := gjson.GetBytes(body2, "instructions").String() require.Equal(t, "existing", existing2) // 测试 3:空白 instructions → 注入 body3 := []byte(`{"model":"gpt-4","instructions":" "}`) existing3 := strings.TrimSpace(gjson.GetBytes(body3, "instructions").String()) require.Empty(t, existing3) // 测试 4:sjson.SetBytes 返回错误时不应 panic // 正常 JSON 不会产生 sjson 错误,验证返回值被正确处理 validBody := []byte(`{"model":"gpt-4"}`) result, setErr := sjson.SetBytes(validBody, "instructions", "hello") require.NoError(t, setErr) require.True(t, gjson.ValidBytes(result)) }