test(bf-413h): add response validation tests and formatting
This commit is contained in:
jedarden
parent
e8c6b4fbb5
commit
4a469b2d52
3 changed files with 668 additions and 1 deletions
|
|
@ -1 +1 @@
|
|||
15ce7dc54ccf06d09bdba2e964f47aaaf7964184
|
||||
3c2482734b9dd8a1bae7cc90c6aaf3f1eff2d091
|
||||
|
|
|
|||
|
|
@ -702,6 +702,10 @@ func TestHelperFunctions(t *testing.T) {
|
|||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Mock Upstream Tests
|
||||
// ============================================================================
|
||||
|
||||
// TestMockUpstreamBasic tests basic mock upstream functionality
|
||||
func TestMockUpstreamBasic(t *testing.T) {
|
||||
t.Run("success scenario", func(t *testing.T) {
|
||||
|
|
|
|||
663
proxy/response_validation_test.go
Normal file
663
proxy/response_validation_test.go
Normal file
|
|
@ -0,0 +1,663 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus/testutil"
|
||||
)
|
||||
|
||||
// TestResponseValidation_EmptyJSONNonStreaming tests that empty JSON body on 2xx responses
|
||||
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
|
||||
func TestResponseValidation_EmptyJSONNonStreaming(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
maxRetries int
|
||||
wantCalls int32 // Expected upstream calls: 1 initial + maxRetries
|
||||
}{
|
||||
{
|
||||
name: "MAX_RETRIES=1",
|
||||
maxRetries: 1,
|
||||
wantCalls: 2, // 1 initial + 1 retry
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=2",
|
||||
maxRetries: 2,
|
||||
wantCalls: 3, // 1 initial + 2 retries
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=5",
|
||||
maxRetries: 5,
|
||||
wantCalls: 6, // 1 initial + 5 retries
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with empty body every time
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Empty body - write nothing
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
start := time.Now()
|
||||
handler.ServeHTTP(w, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
// Verify 502 Bad Gateway is returned after exhausting retries
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify exact upstream call count: 1 initial + maxRetries retries
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != tc.wantCalls {
|
||||
t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
|
||||
tc.wantCalls, tc.maxRetries, actualCalls)
|
||||
}
|
||||
|
||||
// Verify exponential backoff took reasonable time
|
||||
// With maxRetries=N, backoff durations are: 1s, 2s, 4s, ...
|
||||
// For maxRetries=1: ~1s delay
|
||||
// For maxRetries=2: ~1s+2s=3s delay
|
||||
// For maxRetries=5: ~1s+2s+4s+8s+16s=31s delay
|
||||
// We allow 20% tolerance for test execution variance
|
||||
minExpectedDuration := time.Duration(1<<uint(tc.maxRetries-1)) * time.Second
|
||||
if tc.maxRetries > 0 && duration < minExpectedDuration {
|
||||
t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Empty JSON body triggered %d retries, %d upstream calls in %v, returned 502",
|
||||
tc.maxRetries, actualCalls, duration)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponseValidation_InvalidJSONNonStreaming tests that invalid JSON body on 2xx responses
|
||||
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
|
||||
func TestResponseValidation_InvalidJSONNonStreaming(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
maxRetries int
|
||||
wantCalls int32
|
||||
body string
|
||||
}{
|
||||
{
|
||||
name: "MAX_RETRIES=1 with malformed JSON",
|
||||
maxRetries: 1,
|
||||
wantCalls: 2,
|
||||
body: "{invalid json",
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=2 with truncated JSON",
|
||||
maxRetries: 2,
|
||||
wantCalls: 3,
|
||||
body: `{"id": "msg_123", "content":`,
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=3 with garbage data",
|
||||
maxRetries: 3,
|
||||
wantCalls: 4,
|
||||
body: "not json at all!!!",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with invalid JSON every time
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(tc.body))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
start := time.Now()
|
||||
handler.ServeHTTP(w, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
// Verify 502 Bad Gateway is returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify exact upstream call count
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != tc.wantCalls {
|
||||
t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
|
||||
tc.wantCalls, tc.maxRetries, actualCalls)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Invalid JSON body triggered %d retries, %d upstream calls in %v, returned 502",
|
||||
tc.maxRetries, actualCalls, duration)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponseValidation_EmptyStreaming tests that streaming responses opening with zero bytes
|
||||
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
|
||||
func TestResponseValidation_EmptyStreaming(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
maxRetries int
|
||||
wantCalls int32
|
||||
}{
|
||||
{
|
||||
name: "MAX_RETRIES=1",
|
||||
maxRetries: 1,
|
||||
wantCalls: 2,
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=2",
|
||||
maxRetries: 2,
|
||||
wantCalls: 3,
|
||||
},
|
||||
{
|
||||
name: "MAX_RETRIES=4",
|
||||
maxRetries: 4,
|
||||
wantCalls: 5,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with empty streaming response
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Empty streaming - write no data
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
start := time.Now()
|
||||
handler.ServeHTTP(w, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
// Verify 502 Bad Gateway is returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify exact upstream call count
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != tc.wantCalls {
|
||||
t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
|
||||
tc.wantCalls, tc.maxRetries, actualCalls)
|
||||
}
|
||||
|
||||
// Verify exponential backoff took reasonable time
|
||||
minExpectedDuration := time.Duration(1<<uint(tc.maxRetries-1)) * time.Second
|
||||
if tc.maxRetries > 0 && duration < minExpectedDuration {
|
||||
t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Empty streaming triggered %d retries, %d upstream calls in %v, returned 502",
|
||||
tc.maxRetries, actualCalls, duration)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponseValidation_MetricsEmptyJSON tests that zai_proxy_upstream_errors_total
|
||||
// metric increments with correct error_type="truncated_response" for empty JSON retries.
|
||||
func TestResponseValidation_MetricsEmptyJSON(t *testing.T) {
|
||||
// Reset metrics before test
|
||||
retryAttempts.Reset()
|
||||
upstreamErrors.Reset()
|
||||
|
||||
maxRetries := 2
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with empty body
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Empty body
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify zai_proxy_upstream_errors_total metric increments
|
||||
// Each failed attempt (except the final one that returns 502) increments the metric
|
||||
// With maxRetries=2, we have: initial attempt fails (error count 1), retry 1 fails (error count 2), retry 2 fails (error count 3, then returns 502)
|
||||
// So we expect 3 error increments
|
||||
expectedErrorCount := float64(maxRetries + 1)
|
||||
errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test")
|
||||
errorCount := testutil.ToFloat64(errorMetric)
|
||||
if errorCount != expectedErrorCount {
|
||||
t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f",
|
||||
expectedErrorCount, errorCount)
|
||||
}
|
||||
|
||||
// Verify zai_proxy_retry_attempts_total metric increments
|
||||
// Each retry increments the retry attempts metric
|
||||
// With maxRetries=2, we expect 2 retry attempts (not 3 - the final attempt doesn't retry, it returns 502)
|
||||
expectedRetryCount := float64(maxRetries)
|
||||
retryMetric := retryAttempts.WithLabelValues("truncated_response", "test")
|
||||
retryCount := testutil.ToFloat64(retryMetric)
|
||||
if retryCount != expectedRetryCount {
|
||||
t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f",
|
||||
expectedRetryCount, retryCount)
|
||||
}
|
||||
|
||||
// Verify upstream call count
|
||||
expectedCalls := int32(maxRetries + 1)
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != expectedCalls {
|
||||
t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
|
||||
}
|
||||
|
||||
// Verify 502 returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
|
||||
errorCount, retryCount, actualCalls, w.Code)
|
||||
}
|
||||
|
||||
// TestResponseValidation_MetricsInvalidJSON tests that zai_proxy_upstream_errors_total
|
||||
// metric increments with correct error_type="truncated_response" for invalid JSON retries.
|
||||
func TestResponseValidation_MetricsInvalidJSON(t *testing.T) {
|
||||
// Reset metrics before test
|
||||
retryAttempts.Reset()
|
||||
upstreamErrors.Reset()
|
||||
|
||||
maxRetries := 3
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with invalid JSON
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("{this is not valid json"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify upstream errors metric
|
||||
expectedErrorCount := float64(maxRetries + 1)
|
||||
errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test")
|
||||
errorCount := testutil.ToFloat64(errorMetric)
|
||||
if errorCount != expectedErrorCount {
|
||||
t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f",
|
||||
expectedErrorCount, errorCount)
|
||||
}
|
||||
|
||||
// Verify retry attempts metric
|
||||
expectedRetryCount := float64(maxRetries)
|
||||
retryMetric := retryAttempts.WithLabelValues("truncated_response", "test")
|
||||
retryCount := testutil.ToFloat64(retryMetric)
|
||||
if retryCount != expectedRetryCount {
|
||||
t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f",
|
||||
expectedRetryCount, retryCount)
|
||||
}
|
||||
|
||||
// Verify upstream call count
|
||||
expectedCalls := int32(maxRetries + 1)
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != expectedCalls {
|
||||
t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
|
||||
}
|
||||
|
||||
// Verify 502 returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
|
||||
errorCount, retryCount, actualCalls, w.Code)
|
||||
}
|
||||
|
||||
// TestResponseValidation_MetricsEmptyStreaming tests that zai_proxy_upstream_errors_total
|
||||
// metric increments with correct error_type="empty_streaming" for empty streaming retries.
|
||||
func TestResponseValidation_MetricsEmptyStreaming(t *testing.T) {
|
||||
// Reset metrics before test
|
||||
retryAttempts.Reset()
|
||||
upstreamErrors.Reset()
|
||||
|
||||
maxRetries := 2
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
// Create upstream that returns 200 with empty streaming response
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Empty streaming response
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify upstream errors metric with error_type="empty_streaming"
|
||||
expectedErrorCount := float64(maxRetries + 1)
|
||||
errorMetric := upstreamErrors.WithLabelValues("empty_streaming", "test")
|
||||
errorCount := testutil.ToFloat64(errorMetric)
|
||||
if errorCount != expectedErrorCount {
|
||||
t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"empty_streaming\"}} = %f, got %f",
|
||||
expectedErrorCount, errorCount)
|
||||
}
|
||||
|
||||
// Verify retry attempts metric with reason="empty_streaming"
|
||||
expectedRetryCount := float64(maxRetries)
|
||||
retryMetric := retryAttempts.WithLabelValues("empty_streaming", "test")
|
||||
retryCount := testutil.ToFloat64(retryMetric)
|
||||
if retryCount != expectedRetryCount {
|
||||
t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"empty_streaming\"}} = %f, got %f",
|
||||
expectedRetryCount, retryCount)
|
||||
}
|
||||
|
||||
// Verify upstream call count
|
||||
expectedCalls := int32(maxRetries + 1)
|
||||
actualCalls := upstreamCallCount.Load()
|
||||
if actualCalls != expectedCalls {
|
||||
t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
|
||||
}
|
||||
|
||||
// Verify 502 returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
|
||||
t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
|
||||
errorCount, retryCount, actualCalls, w.Code)
|
||||
}
|
||||
|
||||
// TestResponseValidation_ValidResponsePassesThrough tests that valid responses
|
||||
// (both non-streaming and streaming) pass through without triggering retries.
|
||||
func TestResponseValidation_ValidResponsePassesThrough(t *testing.T) {
|
||||
t.Run("valid non-streaming response", func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"id": "msg_test123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": []map[string]string{{"type": "text", "text": "Hello!"}},
|
||||
"usage": map[string]int{"input_tokens": 10, "output_tokens": 5},
|
||||
})
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, 3)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify success response returned
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("Expected status 200 OK, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify response body is valid JSON
|
||||
body := w.Body.String()
|
||||
if !json.Valid([]byte(body)) {
|
||||
t.Error("Response body is not valid JSON")
|
||||
}
|
||||
|
||||
// Verify exactly 1 upstream call (no retries)
|
||||
if upstreamCallCount.Load() != 1 {
|
||||
t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
|
||||
}
|
||||
|
||||
t.Logf("PASS: Valid non-streaming response passed through without retry")
|
||||
})
|
||||
|
||||
t.Run("valid streaming response", func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Send some SSE data
|
||||
w.Write([]byte("data: {\"type\":\"message_start\"}\n\n"))
|
||||
w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n"))
|
||||
w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, 3)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify success response returned
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("Expected status 200 OK, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify exactly 1 upstream call (no retries)
|
||||
if upstreamCallCount.Load() != 1 {
|
||||
t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
|
||||
}
|
||||
|
||||
t.Logf("PASS: Valid streaming response passed through without retry")
|
||||
})
|
||||
}
|
||||
|
||||
// TestResponseValidation_VariousInvalidBodies tests various forms of invalid response bodies
|
||||
// to ensure they all trigger retries correctly.
|
||||
func TestResponseValidation_VariousInvalidBodies(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
body string
|
||||
valid bool
|
||||
}{
|
||||
{
|
||||
name: "empty body",
|
||||
body: "",
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "single opening brace",
|
||||
body: "{",
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "unclosed object",
|
||||
body: `{"key": "value"`,
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "unclosed array",
|
||||
body: `["item1", "item2"`,
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "random text",
|
||||
body: "this is not json at all",
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "partial JSON with trailing comma",
|
||||
body: `{"id": "msg123", "content": "text",}`,
|
||||
valid: false,
|
||||
},
|
||||
{
|
||||
name: "valid JSON object",
|
||||
body: `{"id": "msg123", "content": "hello"}`,
|
||||
valid: true,
|
||||
},
|
||||
{
|
||||
name: "valid JSON array",
|
||||
body: `[{"id": "msg123"}, {"id": "msg456"}]`,
|
||||
valid: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
maxRetries := 1
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
if len(tc.body) > 0 {
|
||||
w.Write([]byte(tc.body))
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
if tc.valid {
|
||||
// Valid JSON should pass through
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("Expected status 200 OK for valid JSON, got %d", w.Code)
|
||||
}
|
||||
if upstreamCallCount.Load() != 1 {
|
||||
t.Errorf("Expected 1 upstream call for valid JSON (no retry), got %d", upstreamCallCount.Load())
|
||||
}
|
||||
} else {
|
||||
// Invalid JSON should trigger retry and return 502
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway for invalid JSON, got %d", w.Code)
|
||||
}
|
||||
expectedCalls := int32(maxRetries + 1)
|
||||
if upstreamCallCount.Load() != expectedCalls {
|
||||
t.Errorf("Expected %d upstream calls for invalid JSON (1 initial + %d retries), got %d",
|
||||
expectedCalls, maxRetries, upstreamCallCount.Load())
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("PASS: %s - status=%d, calls=%d", tc.name, w.Code, upstreamCallCount.Load())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponseValidation_StreamingZeroBytesPeek tests that the streaming zero-byte detection
|
||||
// works by peeking at the first chunk of the response body.
|
||||
func TestResponseValidation_StreamingZeroBytesPeek(t *testing.T) {
|
||||
t.Run("zero bytes on first peek", func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
maxRetries := 1
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Close immediately without writing anything
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, maxRetries)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify retry was triggered and 502 returned
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
|
||||
}
|
||||
expectedCalls := int32(maxRetries + 1)
|
||||
if upstreamCallCount.Load() != expectedCalls {
|
||||
t.Errorf("Expected %d upstream calls, got %d", expectedCalls, upstreamCallCount.Load())
|
||||
}
|
||||
|
||||
t.Logf("PASS: Zero bytes on first peek triggered retry, %d calls, status 502", upstreamCallCount.Load())
|
||||
})
|
||||
|
||||
t.Run("non-zero bytes on first peek (no retry)", func(t *testing.T) {
|
||||
var upstreamCallCount atomic.Int32
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCallCount.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// Write at least one byte immediately
|
||||
w.Write([]byte("d"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
handler := createTestProxyHandler(t, upstream.URL, 3)
|
||||
|
||||
req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
// Verify no retry - response passes through
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("Expected status 200 OK, got %d", w.Code)
|
||||
}
|
||||
if upstreamCallCount.Load() != 1 {
|
||||
t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
|
||||
}
|
||||
|
||||
t.Logf("PASS: Non-zero bytes on first peek, no retry, status 200")
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue