Merge branch 'main' of https://github.com/mt21625457/aicodex2api
This commit is contained in:
yangjianbo
@@ -4,18 +4,42 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// antigravityFailingWriter 模拟客户端断开连接的 gin.ResponseWriter
|
||||
type antigravityFailingWriter struct {
|
||||
gin.ResponseWriter
|
||||
failAfter int // 允许成功写入的次数，之后所有写入返回错误
|
||||
writes int
|
||||
}
|
||||
|
||||
func (w *antigravityFailingWriter) Write(p []byte) (int, error) {
|
||||
if w.writes >= w.failAfter {
|
||||
return 0, errors.New("write failed: client disconnected")
|
||||
}
|
||||
w.writes++
|
||||
return w.ResponseWriter.Write(p)
|
||||
}
|
||||
|
||||
// newAntigravityTestService 创建用于流式测试的 AntigravityGatewayService
|
||||
func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService {
|
||||
return &AntigravityGatewayService{
|
||||
settingService: &SettingService{cfg: cfg},
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) {
|
||||
req := &antigravity.ClaudeRequest{
|
||||
Model: "claude-sonnet-4-5",
|
||||
@@ -338,8 +362,8 @@ func TestAntigravityGatewayService_Forward_StickySessionForceCacheBilling(t *tes
|
||||
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
|
||||
}
|
||||
|
||||
// TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling
|
||||
// 验证：ForwardGemini 粘性会话切换时，UpstreamFailoverError.ForceCacheBilling 应为 true
|
||||
// TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling verifies
|
||||
// that ForwardGemini sets ForceCacheBilling=true for sticky session switch.
|
||||
func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
writer := httptest.NewRecorder()
|
||||
@@ -393,10 +417,16 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(
|
||||
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
|
||||
}
|
||||
|
||||
func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
// TestStreamUpstreamResponse_UsageAndFirstToken
|
||||
// 验证：usage 字段可被累积/覆盖更新，并且能记录首 token 时间
|
||||
func TestStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
@@ -404,25 +434,458 @@ func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte("data: {\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":4}}\n"))
|
||||
_, _ = pw.Write([]byte("data: {\"usage\":{\"output_tokens\":5}}\n"))
|
||||
fmt.Fprintln(pw, `data: {"usage":{"input_tokens":1,"output_tokens":2,"cache_read_input_tokens":3,"cache_creation_input_tokens":4}}`)
|
||||
fmt.Fprintln(pw, `data: {"usage":{"output_tokens":5}}`)
|
||||
}()
|
||||
|
||||
svc := &AntigravityGatewayService{}
|
||||
start := time.Now().Add(-10 * time.Millisecond)
|
||||
usage, firstTokenMs := svc.streamUpstreamResponse(c, resp, start)
|
||||
result := svc.streamUpstreamResponse(c, resp, start)
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, usage)
|
||||
require.Equal(t, 1, usage.InputTokens)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 1, result.usage.InputTokens)
|
||||
// 第二次事件覆盖 output_tokens
|
||||
require.Equal(t, 5, usage.OutputTokens)
|
||||
require.Equal(t, 3, usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, usage.CacheCreationInputTokens)
|
||||
require.Equal(t, 5, result.usage.OutputTokens)
|
||||
require.Equal(t, 3, result.usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, result.usage.CacheCreationInputTokens)
|
||||
require.NotNil(t, result.firstTokenMs)
|
||||
|
||||
if firstTokenMs == nil {
|
||||
t.Fatalf("expected firstTokenMs to be set")
|
||||
}
|
||||
// 确保有透传输出
|
||||
require.True(t, strings.Contains(writer.Body.String(), "data:"))
|
||||
require.Contains(t, rec.Body.String(), "data:")
|
||||
}
|
||||
|
||||
// --- 流式 happy path 测试 ---
|
||||
|
||||
// TestStreamUpstreamResponse_NormalComplete
|
||||
// 验证：正常流式转发完成时，数据正确透传、usage 正确收集、clientDisconnect=false
|
||||
func TestStreamUpstreamResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `event: message_start`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: content_block_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"content_block_delta","delta":{"text":"hello"}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: message_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":5}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 5, result.usage.OutputTokens, "should collect output_tokens from message_delta")
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证数据被透传到客户端
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "event: message_start")
|
||||
require.Contains(t, body, "content_block_delta")
|
||||
require.Contains(t, body, "message_delta")
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_NormalComplete
|
||||
// 验证：正常 Gemini 流式转发，数据正确透传、usage 正确收集
|
||||
func TestHandleGeminiStreamingResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// 第一个 chunk（部分内容）
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
// 第二个 chunk（最终内容+完整 usage）
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":" world"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":8,"cachedContentTokenCount":2}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
// Gemini usage: promptTokenCount=10, candidatesTokenCount=8, cachedContentTokenCount=2
|
||||
// → InputTokens=10-2=8, OutputTokens=8, CacheReadInputTokens=2
|
||||
require.Equal(t, 8, result.usage.InputTokens)
|
||||
require.Equal(t, 8, result.usage.OutputTokens)
|
||||
require.Equal(t, 2, result.usage.CacheReadInputTokens)
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证数据被透传到客户端
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "Hello")
|
||||
require.Contains(t, body, "world")
|
||||
// 不应包含错误事件
|
||||
require.NotContains(t, body, "event: error")
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_NormalComplete
|
||||
// 验证：正常 Claude 流式转发（Gemini→Claude 转换），数据正确转换并输出
|
||||
func TestHandleClaudeStreamingResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// v1internal 包装格式：Gemini 数据嵌套在 "response" 字段下
|
||||
// ProcessLine 先尝试反序列化为 V1InternalResponse，裸格式会导致 Response.UsageMetadata 为空
|
||||
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
// Gemini→Claude 转换的 usage：promptTokenCount=5→InputTokens=5, candidatesTokenCount=3→OutputTokens=3
|
||||
require.Equal(t, 5, result.usage.InputTokens)
|
||||
require.Equal(t, 3, result.usage.OutputTokens)
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证输出是 Claude SSE 格式（processor 会转换）
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "event: message_start", "should contain Claude message_start event")
|
||||
require.Contains(t, body, "event: message_stop", "should contain Claude message_stop event")
|
||||
// 不应包含错误事件
|
||||
require.NotContains(t, body, "event: error")
|
||||
}
|
||||
|
||||
// --- 流式客户端断开检测测试 ---
|
||||
|
||||
// TestStreamUpstreamResponse_ClientDisconnectDrainsUsage
|
||||
// 验证：客户端写入失败后，streamUpstreamResponse 继续读取上游以收集 usage
|
||||
func TestStreamUpstreamResponse_ClientDisconnectDrainsUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `event: message_start`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: message_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":20}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 20, result.usage.OutputTokens)
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_ContextCanceled
|
||||
// 验证：context 取消时返回 usage 且标记 clientDisconnect
|
||||
func TestStreamUpstreamResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_Timeout
|
||||
// 验证：上游超时时返回已收集的 usage
|
||||
func TestStreamUpstreamResponse_Timeout(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pw.Close()
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_TimeoutAfterClientDisconnect
|
||||
// 验证：客户端断开后上游超时，返回 usage 并标记 clientDisconnect
|
||||
func TestStreamUpstreamResponse_TimeoutAfterClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":5}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
// 不关闭 pw → 等待超时
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pw.Close()
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_ClientDisconnect
|
||||
// 验证：Gemini 流式转发中客户端断开后继续 drain 上游
|
||||
func TestHandleGeminiStreamingResponse_ClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":10}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "write_failed")
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_ContextCanceled
|
||||
// 验证：context 取消时不注入错误事件
|
||||
func TestHandleGeminiStreamingResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_ClientDisconnect
|
||||
// 验证：Claude 流式转发中客户端断开后继续 drain 上游
|
||||
func TestHandleClaudeStreamingResponse_ClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// v1internal 包装格式
|
||||
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":15}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_ContextCanceled
|
||||
// 验证：context 取消时不注入错误事件
|
||||
func TestHandleClaudeStreamingResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestExtractSSEUsage 验证 extractSSEUsage 从 SSE data 行正确提取 usage
|
||||
func TestExtractSSEUsage(t *testing.T) {
|
||||
svc := &AntigravityGatewayService{}
|
||||
tests := []struct {
|
||||
name string
|
||||
line string
|
||||
expected ClaudeUsage
|
||||
}{
|
||||
{
|
||||
name: "message_delta with output_tokens",
|
||||
line: `data: {"type":"message_delta","usage":{"output_tokens":42}}`,
|
||||
expected: ClaudeUsage{OutputTokens: 42},
|
||||
},
|
||||
{
|
||||
name: "non-data line ignored",
|
||||
line: `event: message_start`,
|
||||
expected: ClaudeUsage{},
|
||||
},
|
||||
{
|
||||
name: "top-level usage with all fields",
|
||||
line: `data: {"usage":{"input_tokens":10,"output_tokens":20,"cache_read_input_tokens":5,"cache_creation_input_tokens":3}}`,
|
||||
expected: ClaudeUsage{InputTokens: 10, OutputTokens: 20, CacheReadInputTokens: 5, CacheCreationInputTokens: 3},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
usage := &ClaudeUsage{}
|
||||
svc.extractSSEUsage(tt.line, usage)
|
||||
require.Equal(t, tt.expected, *usage)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAntigravityClientWriter 验证 antigravityClientWriter 的断开检测
|
||||
func TestAntigravityClientWriter(t *testing.T) {
|
||||
t.Run("normal write succeeds", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(c.Writer, flusher, "test")
|
||||
|
||||
ok := cw.Write([]byte("hello"))
|
||||
require.True(t, ok)
|
||||
require.False(t, cw.Disconnected())
|
||||
require.Contains(t, rec.Body.String(), "hello")
|
||||
})
|
||||
|
||||
t.Run("write failure marks disconnected", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(fw, flusher, "test")
|
||||
|
||||
ok := cw.Write([]byte("hello"))
|
||||
require.False(t, ok)
|
||||
require.True(t, cw.Disconnected())
|
||||
})
|
||||
|
||||
t.Run("subsequent writes are no-op", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(fw, flusher, "test")
|
||||
|
||||
cw.Write([]byte("first"))
|
||||
ok := cw.Fprintf("second %d", 2)
|
||||
require.False(t, ok)
|
||||
require.True(t, cw.Disconnected())
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user