147 lines
3.9 KiB
Go
147 lines
3.9 KiB
Go
package handler
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type deterministicOpsCaptureWriterStatePool struct {
|
|
states []*opsCaptureWriterState
|
|
}
|
|
|
|
func (p *deterministicOpsCaptureWriterStatePool) Get() any {
|
|
if len(p.states) == 0 {
|
|
return &opsCaptureWriterState{limit: opsCaptureWriterLimit}
|
|
}
|
|
last := len(p.states) - 1
|
|
state := p.states[last]
|
|
p.states = p.states[:last]
|
|
return state
|
|
}
|
|
|
|
func (p *deterministicOpsCaptureWriterStatePool) Put(value any) {
|
|
if state, ok := value.(*opsCaptureWriterState); ok && state != nil {
|
|
p.states = append(p.states, state)
|
|
}
|
|
}
|
|
|
|
func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) {
|
|
w := &opsCaptureWriter{}
|
|
|
|
assert.NotPanics(t, func() {
|
|
assert.Equal(t, 0, w.Status())
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
assert.Equal(t, -1, w.Size())
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
assert.False(t, w.Written())
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
n, err := w.Write([]byte("test"))
|
|
assert.Equal(t, 0, n)
|
|
assert.NoError(t, err)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
n, err := w.WriteString("test")
|
|
assert.Equal(t, 0, n)
|
|
assert.NoError(t, err)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
h := w.Header()
|
|
assert.NotNil(t, h)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
w.WriteHeader(200)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
w.WriteHeaderNow()
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
w.Flush()
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
conn, rw, err := w.Hijack()
|
|
assert.Nil(t, conn)
|
|
assert.Nil(t, rw)
|
|
assert.Error(t, err)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
ch := w.CloseNotify()
|
|
assert.NotNil(t, ch)
|
|
})
|
|
assert.NotPanics(t, func() {
|
|
p := w.Pusher()
|
|
assert.Nil(t, p)
|
|
})
|
|
}
|
|
|
|
func TestOpsCaptureWriter_CompactKeepaliveRestoresOriginalWriter(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
outerStatus := -1
|
|
router.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
outerStatus = c.Writer.Status()
|
|
})
|
|
router.Use(OpsErrorLoggerMiddleware(nil))
|
|
router.GET("/compact", func(c *gin.Context) {
|
|
service.MarkOpenAICompactClientStream(c)
|
|
stop := service.StartOpenAICompactSSEKeepalive(c, time.Hour)
|
|
defer stop()
|
|
c.Status(http.StatusOK)
|
|
})
|
|
|
|
recorder := httptest.NewRecorder()
|
|
request := httptest.NewRequest(http.MethodGet, "/compact", nil)
|
|
require.NotPanics(t, func() {
|
|
router.ServeHTTP(recorder, request)
|
|
})
|
|
require.Equal(t, http.StatusOK, outerStatus)
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
}
|
|
|
|
func TestOpsCaptureWriter_StaleLeaseCannotReachReacquiredState(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
pool := &deterministicOpsCaptureWriterStatePool{}
|
|
|
|
firstRecorder := httptest.NewRecorder()
|
|
firstContext, _ := gin.CreateTestContext(firstRecorder)
|
|
stale := acquireOpsCaptureWriterFromPool(pool, firstContext.Writer)
|
|
releaseOpsCaptureWriter(stale)
|
|
|
|
secondRecorder := httptest.NewRecorder()
|
|
secondContext, _ := gin.CreateTestContext(secondRecorder)
|
|
current := acquireOpsCaptureWriterFromPool(pool, secondContext.Writer)
|
|
defer releaseOpsCaptureWriter(current)
|
|
require.NotSame(t, stale, current)
|
|
require.Same(t, stale.state, current.state)
|
|
|
|
current.WriteHeader(http.StatusInternalServerError)
|
|
_, err := current.WriteString("current")
|
|
require.NoError(t, err)
|
|
require.Equal(t, []byte("current"), current.capturedBytes())
|
|
|
|
n, err := stale.WriteString("stale")
|
|
require.NoError(t, err)
|
|
require.Zero(t, n)
|
|
require.Nil(t, stale.capturedBytes())
|
|
require.Equal(t, []byte("current"), current.capturedBytes())
|
|
require.NotContains(t, secondRecorder.Body.String(), "stale")
|
|
|
|
// Releasing the stale handle must not return an active state to the pool.
|
|
releaseOpsCaptureWriter(stale)
|
|
thirdRecorder := httptest.NewRecorder()
|
|
thirdContext, _ := gin.CreateTestContext(thirdRecorder)
|
|
other := acquireOpsCaptureWriterFromPool(pool, thirdContext.Writer)
|
|
defer releaseOpsCaptureWriter(other)
|
|
require.NotSame(t, current.state, other.state)
|
|
}
|