91 lines
2.0 KiB
Go
91 lines
2.0 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"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) {
|
||
|
|
w := &opsCaptureWriter{}
|
||
|
|
w.ResponseWriter = nil
|
||
|
|
|
||
|
|
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)
|
||
|
|
}
|