e7c24a0c08
Extracted from ardenone-cluster/containers/zai-proxy and ardenone-cluster/containers/zai-proxy-dashboard. - proxy/: OpenAI-compatible ZAI reverse proxy (Go, v1.10.0) - Token counting, rate limiting, Prometheus metrics, canary support - dashboard/: Metrics dashboard backend + React frontend (Go, v1.0.0) - Prometheus collector, SQLite storage, SSE live updates - docs/: Operational notes, research, and plan subdirs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
513 lines
13 KiB
Go
513 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
// TestParseRequestBody tests extracting text from request bodies
|
|
func TestParseRequestBody(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantParts int
|
|
wantEmpty bool
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "Valid single message",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"Hello, how are you?"}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Multiple messages",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"Hi"},{"role":"assistant","content":"Hello there"},{"role":"user","content":"How are you?"}]}`,
|
|
wantParts: 3,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Multi-modal content with text blocks",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":[{"type":"text","text":"Describe this image"}]}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Empty messages array",
|
|
body: `{"model":"glm-4","messages":[]}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Empty message content",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":""}]}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Malformed JSON - invalid syntax",
|
|
body: `{"model": "glm-4", "messages": [}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false, // Graceful degradation
|
|
},
|
|
{
|
|
name: "Malformed JSON - missing messages",
|
|
body: `{"model": "glm-4"}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Malformed JSON - missing model",
|
|
body: `{"messages": [{"role": "user", "content": "hi"}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Malformed JSON - not JSON at all",
|
|
body: `not json at all`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Empty body",
|
|
body: ``,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Long input - multi-turn conversation",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"` + strings.Repeat("The quick brown fox jumps over the lazy dog. ", 100) + `"}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Multi-modal with non-text content",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":[{"type":"image","source":"base64data"},{"type":"text","text":"What is this?"}]}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
wantErr: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
parts, err := ParseRequestBody([]byte(tt.body))
|
|
|
|
if tt.wantErr && err == nil {
|
|
t.Error("Expected error but got none")
|
|
}
|
|
if !tt.wantErr && err != nil {
|
|
t.Errorf("Unexpected error: %v", err)
|
|
}
|
|
|
|
if tt.wantEmpty && len(parts) != 0 {
|
|
t.Errorf("Expected empty result, got %d parts", len(parts))
|
|
}
|
|
|
|
if !tt.wantEmpty && len(parts) != tt.wantParts {
|
|
t.Errorf("ParseRequestBody() returned %d parts, want %d", len(parts), tt.wantParts)
|
|
}
|
|
|
|
// Log extracted text for debugging
|
|
if len(parts) > 0 {
|
|
t.Logf("Extracted %d text parts from request", len(parts))
|
|
for i, part := range parts {
|
|
preview := part
|
|
if len(preview) > 50 {
|
|
preview = preview[:50] + "..."
|
|
}
|
|
t.Logf(" Part %d: %q", i, preview)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestIsStreamingRequest tests detecting streaming requests
|
|
func TestIsStreamingRequest(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantStream bool
|
|
}{
|
|
{
|
|
name: "Streaming enabled",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"Hi"}],"stream":true}`,
|
|
wantStream: true,
|
|
},
|
|
{
|
|
name: "Streaming disabled",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"Hi"}],"stream":false}`,
|
|
wantStream: false,
|
|
},
|
|
{
|
|
name: "Streaming not specified",
|
|
body: `{"model":"glm-4","messages":[{"role":"user","content":"Hi"}]}`,
|
|
wantStream: false,
|
|
},
|
|
{
|
|
name: "Empty body",
|
|
body: ``,
|
|
wantStream: false,
|
|
},
|
|
{
|
|
name: "Invalid JSON",
|
|
body: `{invalid}`,
|
|
wantStream: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := IsStreamingRequest([]byte(tt.body))
|
|
if got != tt.wantStream {
|
|
t.Errorf("IsStreamingRequest() = %v, want %v", got, tt.wantStream)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestBodyTeeReader tests capturing request body while forwarding
|
|
func TestBodyTeeReader(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
content string
|
|
}{
|
|
{
|
|
name: "Simple text",
|
|
content: "Hello, world!",
|
|
},
|
|
{
|
|
name: "JSON request",
|
|
content: `{"model":"glm-4","messages":[{"role":"user","content":"Test"}]}`,
|
|
},
|
|
{
|
|
name: "Long content (10k+ chars)",
|
|
content: strings.Repeat("The quick brown fox jumps over the lazy dog. ", 500),
|
|
},
|
|
{
|
|
name: "Empty content",
|
|
content: "",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
originalBody := io.NopCloser(strings.NewReader(tt.content))
|
|
teeReader := NewBodyTeeReader(originalBody)
|
|
|
|
// Read all content (simulating forwarding to upstream)
|
|
forwarded, err := io.ReadAll(teeReader)
|
|
if err != nil {
|
|
t.Fatalf("Failed to read from TeeReader: %v", err)
|
|
}
|
|
|
|
// Verify forwarded content matches original
|
|
if string(forwarded) != tt.content {
|
|
t.Errorf("Forwarded body mismatch:\ngot: %q\nwant: %q", string(forwarded), tt.content)
|
|
}
|
|
|
|
// Verify captured content matches original
|
|
captured := teeReader.GetCapturedBody()
|
|
if string(captured) != tt.content {
|
|
t.Errorf("Captured body mismatch:\ngot: %q\nwant: %q", string(captured), tt.content)
|
|
}
|
|
|
|
t.Logf("TeeReader successfully captured and forwarded %d bytes", len(tt.content))
|
|
|
|
// Cleanup
|
|
teeReader.Close()
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestBodyTeeReaderConcurrent tests TeeReader doesn't break concurrent reads
|
|
func TestBodyTeeReaderConcurrent(t *testing.T) {
|
|
content := strings.Repeat("test data ", 1000)
|
|
originalBody := io.NopCloser(strings.NewReader(content))
|
|
teeReader := NewBodyTeeReader(originalBody)
|
|
|
|
// Simulate reading in chunks (like proxy forwarding)
|
|
buf := make([]byte, 128)
|
|
var forwarded bytes.Buffer
|
|
|
|
for {
|
|
n, err := teeReader.Read(buf)
|
|
if n > 0 {
|
|
forwarded.Write(buf[:n])
|
|
}
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("Read error: %v", err)
|
|
}
|
|
}
|
|
|
|
// Verify both captured and forwarded are identical
|
|
if forwarded.String() != content {
|
|
t.Errorf("Forwarded content mismatch (length: got %d, want %d)", forwarded.Len(), len(content))
|
|
}
|
|
|
|
captured := string(teeReader.GetCapturedBody())
|
|
if captured != content {
|
|
t.Errorf("Captured content mismatch (length: got %d, want %d)", len(captured), len(content))
|
|
}
|
|
|
|
teeReader.Close()
|
|
}
|
|
|
|
// TestParseResponseBody tests extracting text from non-streaming responses
|
|
func TestParseResponseBody(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantParts int
|
|
wantEmpty bool
|
|
}{
|
|
{
|
|
name: "Valid response with single content block",
|
|
body: `{"id":"msg_123","type":"message","role":"assistant","content":[{"type":"text","text":"Hello! How can I help you?"}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
},
|
|
{
|
|
name: "Response with multiple content blocks",
|
|
body: `{"content":[{"type":"text","text":"First part"},{"type":"text","text":"Second part"}]}`,
|
|
wantParts: 2,
|
|
wantEmpty: false,
|
|
},
|
|
{
|
|
name: "Empty content array",
|
|
body: `{"content":[]}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
},
|
|
{
|
|
name: "Invalid JSON",
|
|
body: `{invalid}`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
},
|
|
{
|
|
name: "Empty body",
|
|
body: ``,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
},
|
|
{
|
|
name: "Long response",
|
|
body: `{"content":[{"type":"text","text":"` + strings.Repeat("word ", 5000) + `"}]}`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
parts, err := ParseResponseBody([]byte(tt.body))
|
|
if err != nil {
|
|
t.Errorf("ParseResponseBody() error = %v", err)
|
|
return
|
|
}
|
|
|
|
if tt.wantEmpty && len(parts) != 0 {
|
|
t.Errorf("Expected empty result, got %d parts", len(parts))
|
|
}
|
|
|
|
if !tt.wantEmpty && len(parts) != tt.wantParts {
|
|
t.Errorf("ParseResponseBody() returned %d parts, want %d", len(parts), tt.wantParts)
|
|
}
|
|
|
|
if len(parts) > 0 {
|
|
t.Logf("Extracted %d text parts from response", len(parts))
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestParseSSEResponse tests extracting text from streaming responses
|
|
func TestParseSSEResponse(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantParts int
|
|
wantEmpty bool
|
|
}{
|
|
{
|
|
name: "Valid SSE stream with multiple deltas",
|
|
body: `data: {"type":"message_start","message":{"id":"msg_123"}}
|
|
|
|
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
|
|
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
|
|
|
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" there"}}
|
|
|
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"!"}}
|
|
|
|
data: {"type":"content_block_stop","index":0}
|
|
|
|
data: {"type":"message_stop"}
|
|
`,
|
|
wantParts: 3,
|
|
wantEmpty: false,
|
|
},
|
|
{
|
|
name: "SSE with single delta",
|
|
body: `data: {"type":"content_block_delta","index":0,"delta":{"text":"Hello world"}}
|
|
|
|
`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
},
|
|
{
|
|
name: "SSE with no deltas",
|
|
body: `data: {"type":"message_start"}
|
|
|
|
data: {"type":"message_stop"}
|
|
`,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
},
|
|
{
|
|
name: "Empty SSE",
|
|
body: ``,
|
|
wantParts: 0,
|
|
wantEmpty: true,
|
|
},
|
|
{
|
|
name: "Malformed SSE - invalid JSON",
|
|
body: `data: {invalid json}
|
|
|
|
data: {"type":"content_block_delta","delta":{"text":"Valid"}}
|
|
`,
|
|
wantParts: 1,
|
|
wantEmpty: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
parts, err := ParseSSEResponse([]byte(tt.body))
|
|
if err != nil {
|
|
t.Errorf("ParseSSEResponse() error = %v", err)
|
|
return
|
|
}
|
|
|
|
if tt.wantEmpty && len(parts) != 0 {
|
|
t.Errorf("Expected empty result, got %d parts", len(parts))
|
|
}
|
|
|
|
if !tt.wantEmpty && len(parts) != tt.wantParts {
|
|
t.Errorf("ParseSSEResponse() returned %d parts, want %d", len(parts), tt.wantParts)
|
|
}
|
|
|
|
if len(parts) > 0 {
|
|
t.Logf("Extracted %d text deltas from SSE response", len(parts))
|
|
for i, part := range parts {
|
|
t.Logf(" Delta %d: %q", i, part)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestIsSSEFormat tests SSE format detection
|
|
func TestIsSSEFormat(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantSSE bool
|
|
}{
|
|
{
|
|
name: "Valid SSE",
|
|
body: `data: {"type":"message_start"}`,
|
|
wantSSE: true,
|
|
},
|
|
{
|
|
name: "JSON response",
|
|
body: `{"type":"message","content":[]}`,
|
|
wantSSE: false,
|
|
},
|
|
{
|
|
name: "Empty body",
|
|
body: ``,
|
|
wantSSE: false,
|
|
},
|
|
{
|
|
name: "Plain text with 'data:' in content",
|
|
body: `This is plain text that mentions data: fields`,
|
|
wantSSE: true, // Contains "data: " pattern
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := IsSSEFormat([]byte(tt.body))
|
|
if got != tt.wantSSE {
|
|
t.Errorf("IsSSEFormat() = %v, want %v", got, tt.wantSSE)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// BenchmarkParseRequestBody benchmarks request body parsing
|
|
func BenchmarkParseRequestBody(b *testing.B) {
|
|
body := []byte(`{"model":"glm-4","messages":[{"role":"user","content":"Hello, how are you today?"}]}`)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
ParseRequestBody(body)
|
|
}
|
|
}
|
|
|
|
// BenchmarkParseRequestBodyLong benchmarks parsing with long content
|
|
func BenchmarkParseRequestBodyLong(b *testing.B) {
|
|
longText := strings.Repeat("The quick brown fox jumps over the lazy dog. ", 500)
|
|
body := []byte(`{"model":"glm-4","messages":[{"role":"user","content":"` + longText + `"}]}`)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
ParseRequestBody(body)
|
|
}
|
|
}
|
|
|
|
// BenchmarkBodyTeeReader benchmarks TeeReader performance
|
|
func BenchmarkBodyTeeReader(b *testing.B) {
|
|
content := strings.Repeat("test data ", 100)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
originalBody := io.NopCloser(strings.NewReader(content))
|
|
teeReader := NewBodyTeeReader(originalBody)
|
|
io.ReadAll(teeReader)
|
|
teeReader.Close()
|
|
}
|
|
}
|
|
|
|
// BenchmarkParseSSEResponse benchmarks SSE parsing
|
|
func BenchmarkParseSSEResponse(b *testing.B) {
|
|
body := []byte(`data: {"type":"content_block_delta","delta":{"text":"Hello"}}
|
|
|
|
data: {"type":"content_block_delta","delta":{"text":" world"}}
|
|
`)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
ParseSSEResponse(body)
|
|
}
|
|
}
|