78 lines
2.1 KiB
Go
78 lines
2.1 KiB
Go
|
package server
|
||
|
|
|
||
|
|
import (
|
||
|
|
"encoding/json"
|
||
|
|
"net/http/httptest"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/daniel/cursor-adapter/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestNewSSEWriter(t *testing.T) {
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
sse := NewSSEWriter(rec)
|
||
|
|
if sse == nil {
|
||
|
|
t.Fatal("NewSSEWriter returned nil")
|
||
|
|
}
|
||
|
|
headers := rec.Header()
|
||
|
|
if got := headers.Get("Content-Type"); got != "text/event-stream" {
|
||
|
|
t.Errorf("Content-Type = %q, want %q", got, "text/event-stream")
|
||
|
|
}
|
||
|
|
if got := headers.Get("Cache-Control"); got != "no-cache" {
|
||
|
|
t.Errorf("Cache-Control = %q, want %q", got, "no-cache")
|
||
|
|
}
|
||
|
|
if got := headers.Get("Connection"); got != "keep-alive" {
|
||
|
|
t.Errorf("Connection = %q, want %q", got, "keep-alive")
|
||
|
|
}
|
||
|
|
if got := headers.Get("X-Accel-Buffering"); got != "no" {
|
||
|
|
t.Errorf("X-Accel-Buffering = %q, want %q", got, "no")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWriteChunk(t *testing.T) {
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
sse := NewSSEWriter(rec)
|
||
|
|
|
||
|
|
content := "hello"
|
||
|
|
chunk := types.NewChatCompletionChunk("test-id", 0, "", types.Delta{Content: &content})
|
||
|
|
|
||
|
|
if err := sse.WriteChunk(chunk); err != nil {
|
||
|
|
t.Fatalf("WriteChunk returned error: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
body := rec.Body.String()
|
||
|
|
if !strings.HasPrefix(body, "data: ") {
|
||
|
|
t.Errorf("WriteChunk output missing 'data: ' prefix, got %q", body)
|
||
|
|
}
|
||
|
|
if !strings.HasSuffix(body, "\n\n") {
|
||
|
|
t.Errorf("WriteChunk output missing trailing newlines, got %q", body)
|
||
|
|
}
|
||
|
|
|
||
|
|
jsonStr := strings.TrimPrefix(body, "data: ")
|
||
|
|
jsonStr = strings.TrimSuffix(jsonStr, "\n\n")
|
||
|
|
var parsed types.ChatCompletionChunk
|
||
|
|
if err := json.Unmarshal([]byte(jsonStr), &parsed); err != nil {
|
||
|
|
t.Fatalf("failed to unmarshal chunk JSON: %v", err)
|
||
|
|
}
|
||
|
|
if parsed.ID != "test-id" {
|
||
|
|
t.Errorf("parsed chunk ID = %q, want %q", parsed.ID, "test-id")
|
||
|
|
}
|
||
|
|
if len(parsed.Choices) != 1 || *parsed.Choices[0].Delta.Content != "hello" {
|
||
|
|
t.Errorf("parsed chunk content mismatch, got %v", parsed)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWriteDone(t *testing.T) {
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
sse := NewSSEWriter(rec)
|
||
|
|
|
||
|
|
sse.WriteDone()
|
||
|
|
|
||
|
|
body := rec.Body.String()
|
||
|
|
want := "data: [DONE]\n\n"
|
||
|
|
if body != want {
|
||
|
|
t.Errorf("WriteDone output = %q, want %q", body, want)
|
||
|
|
}
|
||
|
|
}
|