134 lines
3.6 KiB
Go
134 lines
3.6 KiB
Go
package v1
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"context"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/gochat/gochat/internal/llm"
|
|
"github.com/gochat/gochat/internal/model"
|
|
)
|
|
|
|
// ========== Mock LLM Provider for SSE tests ==========
|
|
|
|
type mockSSELLMProvider struct {
|
|
streamChunks []llm.StreamChunk
|
|
streamError error
|
|
lastStreamReq *llm.ChatRequest
|
|
}
|
|
|
|
func (m *mockSSELLMProvider) ChatCompletion(ctx context.Context, req llm.ChatRequest) (*llm.ChatResponse, error) {
|
|
return nil, fmt.Errorf("not implemented")
|
|
}
|
|
|
|
func (m *mockSSELLMProvider) CreateEmbedding(ctx context.Context, req llm.EmbeddingRequest) (*llm.EmbeddingResponse, error) {
|
|
return nil, fmt.Errorf("not implemented")
|
|
}
|
|
|
|
func (m *mockSSELLMProvider) ChatCompletionStream(ctx context.Context, req llm.ChatRequest, onChunk func(llm.StreamChunk) error) error {
|
|
m.lastStreamReq = &req
|
|
if m.streamError != nil {
|
|
return m.streamError
|
|
}
|
|
for _, chunk := range m.streamChunks {
|
|
if err := onChunk(chunk); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ========== escapeJSONString Tests ==========
|
|
|
|
func TestEscapeJSONString_双引号(t *testing.T) {
|
|
result := escapeJSONString(`he said "hello"`)
|
|
assert.Equal(t, `he said \"hello\"`, result)
|
|
}
|
|
|
|
func TestEscapeJSONString_反斜杠(t *testing.T) {
|
|
result := escapeJSONString(`path\to\file`)
|
|
assert.Equal(t, `path\\to\\file`, result)
|
|
}
|
|
|
|
func TestEscapeJSONString_换行(t *testing.T) {
|
|
result := escapeJSONString("line1\nline2")
|
|
assert.Equal(t, "line1\\nline2", result)
|
|
}
|
|
|
|
func TestEscapeJSONString_制表符(t *testing.T) {
|
|
result := escapeJSONString("col1\tcol2")
|
|
assert.Equal(t, "col1\\tcol2", result)
|
|
}
|
|
|
|
func TestEscapeJSONString_无特殊字符(t *testing.T) {
|
|
result := escapeJSONString("plain text")
|
|
assert.Equal(t, "plain text", result)
|
|
}
|
|
|
|
// ========== buildStreamChatMessages Tests ==========
|
|
|
|
func TestBuildStreamChatMessages_空历史(t *testing.T) {
|
|
thread := &model.CopilotThread{
|
|
Title: "Test Thread",
|
|
}
|
|
_ = thread // used indirectly via buildStreamChatMessages
|
|
|
|
result := buildStreamChatMessages(thread, "user question")
|
|
|
|
assert.Len(t, result, 2) // system + user
|
|
assert.Equal(t, "system", result[0].Role)
|
|
assert.Equal(t, "user", result[1].Role)
|
|
assert.Equal(t, "user question", result[1].Content)
|
|
}
|
|
|
|
func TestBuildStreamChatMessages_有历史(t *testing.T) {
|
|
thread := &model.CopilotThread{
|
|
Title: "Test Thread",
|
|
}
|
|
messages := []model.CopilotMessage{
|
|
{MessageType: model.CopilotMessageTypeUser, Message: json.RawMessage(`"previous user msg"`)},
|
|
{MessageType: model.CopilotMessageTypeAssistant, Message: json.RawMessage(`"previous assistant msg"`)},
|
|
}
|
|
thread.Messages = messages
|
|
|
|
result := buildStreamChatMessages(thread, "new question")
|
|
|
|
// system + 2 history + new user = 4
|
|
assert.Len(t, result, 4)
|
|
assert.Equal(t, "system", result[0].Role)
|
|
assert.Equal(t, "new question", result[len(result)-1].Content)
|
|
}
|
|
|
|
// ========== writeSSEMessage Tests ==========
|
|
|
|
func TestWriteSSEMessage_格式(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
|
|
writeSSEMessage(c, "message", `{"content": "hello"}`)
|
|
|
|
body := w.Body.String()
|
|
assert.Contains(t, body, "event: message")
|
|
assert.Contains(t, body, `data: {"content": "hello"}`)
|
|
// SSE format requires double newline
|
|
assert.Contains(t, body, "\n\n")
|
|
}
|
|
|
|
func TestWriteSSEMessage_Done事件(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
|
|
writeSSEMessage(c, "done", `{"done": true}`)
|
|
|
|
body := w.Body.String()
|
|
assert.Contains(t, body, "event: done")
|
|
assert.Contains(t, body, `data: {"done": true}`)
|
|
}
|