//go:build unit package service import ( "context" "io" "net/http" "net/http/httptest" "testing" "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) // --- parseSSEUsage 测试 --- func newMinimalGatewayService() *GatewayService { return &GatewayService{ cfg: &config.Config{ Gateway: config.GatewayConfig{ StreamDataIntervalTimeout: 0, MaxLineSize: defaultMaxLineSize, }, }, rateLimitService: &RateLimitService{}, } } func TestParseSSEUsage_MessageStart(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} data := `{"type":"message_start","message":{"usage":{"input_tokens":100,"cache_creation_input_tokens":50,"cache_read_input_tokens":200}}}` svc.parseSSEUsage(data, usage) require.Equal(t, 100, usage.InputTokens) require.Equal(t, 50, usage.CacheCreationInputTokens) require.Equal(t, 200, usage.CacheReadInputTokens) require.Equal(t, 0, usage.OutputTokens, "message_start 不应设置 output_tokens") } func TestParseSSEUsage_MessageDelta(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} data := `{"type":"message_delta","usage":{"output_tokens":42}}` svc.parseSSEUsage(data, usage) require.Equal(t, 42, usage.OutputTokens) require.Equal(t, 0, usage.InputTokens, "message_delta 的 output_tokens 不应影响已有的 input_tokens") } func TestParseSSEUsage_DeltaDoesNotOverwriteStartValues(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} // 先处理 message_start svc.parseSSEUsage(`{"type":"message_start","message":{"usage":{"input_tokens":100}}}`, usage) require.Equal(t, 100, usage.InputTokens) // 再处理 message_delta(output_tokens > 0, input_tokens = 0) svc.parseSSEUsage(`{"type":"message_delta","usage":{"output_tokens":50}}`, usage) require.Equal(t, 100, usage.InputTokens, "delta 中 input_tokens=0 不应覆盖 start 中的值") require.Equal(t, 50, usage.OutputTokens) } func TestParseSSEUsage_DeltaOverwritesWithNonZero(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} // GLM 等 API 会在 delta 中包含所有 usage 信息 svc.parseSSEUsage(`{"type":"message_delta","usage":{"input_tokens":200,"output_tokens":100,"cache_creation_input_tokens":30,"cache_read_input_tokens":60}}`, usage) require.Equal(t, 200, usage.InputTokens) require.Equal(t, 100, usage.OutputTokens) require.Equal(t, 30, usage.CacheCreationInputTokens) require.Equal(t, 60, usage.CacheReadInputTokens) } func TestParseSSEUsage_InvalidJSON(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} // 无效 JSON 不应 panic svc.parseSSEUsage("not json", usage) require.Equal(t, 0, usage.InputTokens) require.Equal(t, 0, usage.OutputTokens) } func TestParseSSEUsage_UnknownType(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} // 不是 message_start 或 message_delta 的类型 svc.parseSSEUsage(`{"type":"content_block_delta","delta":{"text":"hello"}}`, usage) require.Equal(t, 0, usage.InputTokens) require.Equal(t, 0, usage.OutputTokens) } func TestParseSSEUsage_EmptyString(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} svc.parseSSEUsage("", usage) require.Equal(t, 0, usage.InputTokens) } func TestParseSSEUsage_DoneEvent(t *testing.T) { svc := newMinimalGatewayService() usage := &ClaudeUsage{} // [DONE] 事件不应影响 usage svc.parseSSEUsage("[DONE]", usage) require.Equal(t, 0, usage.InputTokens) } // --- 流式响应端到端测试 --- func TestHandleStreamingResponse_CacheTokens(t *testing.T) { gin.SetMode(gin.TestMode) svc := newMinimalGatewayService() rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) pr, pw := io.Pipe() resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: pr} go func() { defer func() { _ = pw.Close() }() _, _ = pw.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10,\"cache_creation_input_tokens\":20,\"cache_read_input_tokens\":30}}}\n\n")) _, _ = pw.Write([]byte("data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":15}}\n\n")) _, _ = pw.Write([]byte("data: [DONE]\n\n")) }() result, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "model", "model", false) _ = pr.Close() require.NoError(t, err) require.NotNil(t, result) require.NotNil(t, result.usage) require.Equal(t, 10, result.usage.InputTokens) require.Equal(t, 15, result.usage.OutputTokens) require.Equal(t, 20, result.usage.CacheCreationInputTokens) require.Equal(t, 30, result.usage.CacheReadInputTokens) } func TestHandleStreamingResponse_EmptyStream(t *testing.T) { gin.SetMode(gin.TestMode) svc := newMinimalGatewayService() rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) pr, pw := io.Pipe() resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: pr} go func() { // 直接关闭,不发送任何事件 _ = pw.Close() }() result, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "model", "model", false) _ = pr.Close() require.NoError(t, err) require.NotNil(t, result) } func TestHandleStreamingResponse_SpecialCharactersInJSON(t *testing.T) { gin.SetMode(gin.TestMode) svc := newMinimalGatewayService() rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) pr, pw := io.Pipe() resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: pr} go func() { defer func() { _ = pw.Close() }() // 包含特殊字符的 content_block_delta(引号、换行、Unicode) _, _ = pw.Write([]byte("data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello \\\"world\\\"\\n你好\"}}\n\n")) _, _ = pw.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":5}}}\n\n")) _, _ = pw.Write([]byte("data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":3}}\n\n")) _, _ = pw.Write([]byte("data: [DONE]\n\n")) }() result, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "model", "model", false) _ = pr.Close() require.NoError(t, err) require.NotNil(t, result) require.NotNil(t, result.usage) require.Equal(t, 5, result.usage.InputTokens) require.Equal(t, 3, result.usage.OutputTokens) // 验证响应中包含转发的数据 body := rec.Body.String() require.Contains(t, body, "content_block_delta", "响应应包含转发的 SSE 事件") }