836ba14b70
- handleUpstreamError 补齐新增的三个参数 (0, "", false) - handleStreamingResponse 移除已删除的 nil 参数 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
53 lines
1.5 KiB
Go
53 lines
1.5 KiB
Go
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"
|
|
)
|
|
|
|
func TestGatewayService_StreamingReusesScannerBufferAndStillParsesUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
cfg := &config.Config{
|
|
Gateway: config.GatewayConfig{
|
|
StreamDataIntervalTimeout: 0,
|
|
MaxLineSize: defaultMaxLineSize,
|
|
},
|
|
}
|
|
|
|
svc := &GatewayService{
|
|
cfg: cfg,
|
|
rateLimitService: &RateLimitService{},
|
|
}
|
|
|
|
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() }()
|
|
// Minimal SSE event to trigger parseSSEUsage
|
|
_, _ = pw.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":3}}}\n\n"))
|
|
_, _ = pw.Write([]byte("data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":7}}\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, 3, result.usage.InputTokens)
|
|
require.Equal(t, 7, result.usage.OutputTokens)
|
|
}
|