refactor: extract ReadUpstreamResponseBody to deduplicate upstream response read + too-large error handling
Consolidates 9 call sites of resolveUpstreamResponseReadLimit + readUpstreamResponseBodyLimited + ErrUpstreamResponseBodyTooLarge error handling into a single ReadUpstreamResponseBody function with TooLargeWriter callback for API-format-specific error responses (Anthropic, OpenAI, countTokens).
This commit is contained in:
@@ -4,8 +4,10 @@ import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"testing"
|
||||
"testing/iotest"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -35,3 +37,44 @@ func TestReadUpstreamResponseBodyLimited(t *testing.T) {
|
||||
require.True(t, errors.Is(err, ErrUpstreamResponseBodyTooLarge))
|
||||
})
|
||||
}
|
||||
|
||||
func TestReadUpstreamResponseBody(t *testing.T) {
|
||||
t.Run("within limit", func(t *testing.T) {
|
||||
body, err := ReadUpstreamResponseBody(bytes.NewReader([]byte("ok")), nil, nil, nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("ok"), body)
|
||||
})
|
||||
|
||||
t.Run("exceeds limit calls onTooLarge", func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.UpstreamResponseReadMaxBytes = 3
|
||||
|
||||
called := false
|
||||
onTooLarge := func(_ *gin.Context) { called = true }
|
||||
|
||||
body, err := ReadUpstreamResponseBody(bytes.NewReader([]byte("toolong")), cfg, nil, onTooLarge)
|
||||
require.Nil(t, body)
|
||||
require.True(t, errors.Is(err, ErrUpstreamResponseBodyTooLarge))
|
||||
require.True(t, called)
|
||||
})
|
||||
|
||||
t.Run("nil onTooLarge does not panic", func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.UpstreamResponseReadMaxBytes = 3
|
||||
|
||||
body, err := ReadUpstreamResponseBody(bytes.NewReader([]byte("toolong")), cfg, nil, nil)
|
||||
require.Nil(t, body)
|
||||
require.True(t, errors.Is(err, ErrUpstreamResponseBodyTooLarge))
|
||||
})
|
||||
|
||||
t.Run("io error does not call onTooLarge", func(t *testing.T) {
|
||||
called := false
|
||||
onTooLarge := func(_ *gin.Context) { called = true }
|
||||
|
||||
body, err := ReadUpstreamResponseBody(iotest.ErrReader(errors.New("disk failure")), nil, nil, onTooLarge)
|
||||
require.Nil(t, body)
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.Is(err, ErrUpstreamResponseBodyTooLarge))
|
||||
require.False(t, called)
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user