sub2api/backend/internal/handler/ops_capture_writer_nil_test.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 blockingOpsResponseWriter struct {
	gin.ResponseWriter
	writeStarted chan struct{}
	writeRelease chan struct{}
}

func (w *blockingOpsResponseWriter) WriteString(s string) (int, error) {
	close(w.writeStarted)
	<-w.writeRelease
	return w.ResponseWriter.WriteString(s)
}

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)
}

func TestOpsCaptureWriter_ReleaseWaitsForDelegatedWriteWithoutHoldingStateMutex(t *testing.T) {
	gin.SetMode(gin.TestMode)
	pool := &deterministicOpsCaptureWriterStatePool{}
	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)
	inner := &blockingOpsResponseWriter{
		ResponseWriter: ctx.Writer,
		writeStarted:   make(chan struct{}),
		writeRelease:   make(chan struct{}),
	}
	w := acquireOpsCaptureWriterFromPool(pool, inner)

	writeDone := make(chan struct{})
	go func() {
		defer close(writeDone)
		_, _ = w.WriteString("body")
	}()
	<-inner.writeStarted

	if !w.state.mu.TryLock() {
		t.Fatal("state mutex remained held across the delegated network write")
	}
	w.state.mu.Unlock()

	releaseDone := make(chan struct{})
	go func() {
		releaseOpsCaptureWriter(w)
		close(releaseDone)
	}()
	select {
	case <-releaseDone:
		t.Fatal("release returned while a delegated write was still active")
	case <-time.After(20 * time.Millisecond):
	}
	require.Empty(t, pool.states)

	close(inner.writeRelease)
	<-writeDone
	select {
	case <-releaseDone:
	case <-time.After(time.Second):
		t.Fatal("release did not finish after the delegated write returned")
	}
	require.Len(t, pool.states, 1)
}