修复代理与日志链路的可靠性问题
This commit is contained in:
@@ -1,7 +1,9 @@
|
||||
package config_test
|
||||
|
||||
import (
|
||||
"math"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -50,6 +52,9 @@ func TestLoadQueueDefaults(t *testing.T) {
|
||||
if cfg.MaxBodyBytes != 1<<20 {
|
||||
t.Fatalf("MaxBodyBytes=%d want %d", cfg.MaxBodyBytes, 1<<20)
|
||||
}
|
||||
if cfg.MaxRequestBytes != 16<<20 {
|
||||
t.Fatalf("MaxRequestBytes=%d want %d", cfg.MaxRequestBytes, 16<<20)
|
||||
}
|
||||
if cfg.LogQueueSize != 256 {
|
||||
t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize)
|
||||
}
|
||||
@@ -242,6 +247,23 @@ func TestLoadRequiresClickHouseURL(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsUnsupportedUpstreamScheme(t *testing.T) {
|
||||
t.Setenv("UPSTREAM_URL", "ftp://example.com")
|
||||
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
|
||||
if _, err := config.Load(); err == nil {
|
||||
t.Fatal("Load succeeded with unsupported UPSTREAM_URL scheme")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsMaxRequestBytesOverflowBoundary(t *testing.T) {
|
||||
t.Setenv("UPSTREAM_URL", "https://example.com")
|
||||
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
|
||||
t.Setenv("MAX_REQUEST_BYTES", strconv.FormatInt(math.MaxInt64, 10))
|
||||
if _, err := config.Load(); err == nil || !strings.Contains(err.Error(), "MAX_REQUEST_BYTES") {
|
||||
t.Fatalf("Load error=%v, want MAX_REQUEST_BYTES rejection", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadValidatesClickHouseURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -244,8 +244,8 @@ func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
|
||||
RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"),
|
||||
ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"),
|
||||
}
|
||||
if got, want := logger.EstimatedBytes(entry), int64(55); got != want {
|
||||
t.Fatalf("EstimatedBytes=%d want %d", got, want)
|
||||
if got := logger.EstimatedBytes(entry); got <= 55 {
|
||||
t.Fatalf("EstimatedBytes=%d must include fixed entry overhead", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -485,6 +485,7 @@ func TestConcurrentStopCannotCancelFirstStopDrain(t *testing.T) {
|
||||
firstCtx, cancelFirst := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelFirst()
|
||||
go func() { firstResult <- q.Stop(firstCtx) }()
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
for q.Stats().Dropped == 0 {
|
||||
q.Submit(&logger.LogEntry{RequestID: "stop-probe"})
|
||||
}
|
||||
|
||||
@@ -108,6 +108,38 @@ func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestBodyOverLimitReturns413WithoutUpstream(t *testing.T) {
|
||||
var upstreamCalls atomic.Int32
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCalls.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
u, _ := url.Parse(upstream.URL)
|
||||
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{MaxRequestBytes: 4})
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
contentLength int64
|
||||
}{
|
||||
{name: "known length", contentLength: 5},
|
||||
{name: "chunked", contentLength: -1},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", strings.NewReader("12345"))
|
||||
req.ContentLength = tc.contentLength
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusRequestEntityTooLarge || rec.Body.String() != "request body too large\n" {
|
||||
t.Fatalf("response=(%d, %q), want fixed 413", rec.Code, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
if upstreamCalls.Load() != 0 {
|
||||
t.Fatalf("upstream called %d times", upstreamCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestBodyReadFailureWithZeroContentLengthDoesNotReachUpstream(t *testing.T) {
|
||||
var upstreamCalls atomic.Int32
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -643,6 +643,58 @@ func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiStreamPreservesMultiplePartsByIndex(t *testing.T) {
|
||||
upstream := newSSEUpstream([]string{
|
||||
`data: {"candidates":[{"content":{"parts":[{"text":"hello "},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"world"}],"role":"model"},"index":0}]}` + "\n\n",
|
||||
`data: {"candidates":[{"content":{"parts":[{"text":"again"},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"!"}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
|
||||
})
|
||||
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
|
||||
|
||||
var captured struct {
|
||||
Candidates []struct {
|
||||
Content struct {
|
||||
Parts []map[string]any `json:"parts"`
|
||||
} `json:"content"`
|
||||
} `json:"candidates"`
|
||||
}
|
||||
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
|
||||
t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
|
||||
}
|
||||
parts := captured.Candidates[0].Content.Parts
|
||||
if len(parts) != 3 || parts[0]["text"] != "hello again" || parts[2]["text"] != "world!" {
|
||||
t.Fatalf("multiple Gemini parts not preserved: %+v", parts)
|
||||
}
|
||||
if _, ok := parts[1]["functionCall"]; !ok {
|
||||
t.Fatalf("middle functionCall part missing: %+v", parts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiStreamAppendsDifferentPartKindsAcrossChunks(t *testing.T) {
|
||||
upstream := newSSEUpstream([]string{
|
||||
`data: {"candidates":[{"content":{"parts":[{"text":"answer"}],"role":"model"},"index":0}]}` + "\n\n",
|
||||
`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
|
||||
})
|
||||
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
|
||||
|
||||
var captured struct {
|
||||
Candidates []struct {
|
||||
Content struct {
|
||||
Parts []map[string]any `json:"parts"`
|
||||
} `json:"content"`
|
||||
} `json:"candidates"`
|
||||
}
|
||||
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
|
||||
t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
|
||||
}
|
||||
parts := captured.Candidates[0].Content.Parts
|
||||
if len(parts) != 2 || parts[0]["text"] != "answer" {
|
||||
t.Fatalf("different Gemini part kinds were merged: %+v", parts)
|
||||
}
|
||||
if _, ok := parts[1]["functionCall"]; !ok {
|
||||
t.Fatalf("functionCall part missing: %+v", parts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownSSEKeepsRawBody(t *testing.T) {
|
||||
e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user