sub2api/backend/internal/handler/gemini_client_cancel_test.go

//go:build unit

package handler

import (
	"bytes"
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sync/atomic"
	"testing"

	"github.com/Wei-Shaw/sub2api/internal/config"
	"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
	middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

// clientCancelUpstream 模拟上游尚未响应时客户端断开：Do 期间取消请求 context，
// 并按 net/http 的形态返回包裹 context.Canceled 的 *url.Error。
type clientCancelUpstream struct {
	service.HTTPUpstream
	cancel context.CancelFunc
	calls  atomic.Int32
}

func (u *clientCancelUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
	u.calls.Add(1)
	u.cancel()
	return nil, &url.Error{Op: "Post", URL: req.URL.String(), Err: context.Canceled}
}

type geminiClientCancelFixture struct {
	handler  *GatewayHandler
	group    *service.Group
	apiKey   *service.APIKey
	upstream *clientCancelUpstream
	ctx      context.Context
}

func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture {
	t.Helper()
	groupID := int64(9200)
	accountID := int64(9201)
	group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive}
	account := &service.Account{
		ID:            accountID,
		Platform:      service.PlatformGemini,
		Type:          service.AccountTypeAPIKey,
		Credentials:   map[string]any{"api_key": "test-key"},
		Status:        service.StatusActive,
		Schedulable:   true,
		Concurrency:   1,
		Priority:      1,
		AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}},
	}
	h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account})
	t.Cleanup(cleanup)

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	upstream := &clientCancelUpstream{cancel: cancel}
	h.geminiCompatService = service.NewGeminiMessagesCompatService(nil, nil, nil, nil, nil, nil, upstream, nil, &config.Config{})

	apiKey := &service.APIKey{
		ID: 9202, UserID: 9203, GroupID: &groupID, Group: group, Status: service.StatusActive,
		User: &service.User{ID: 9203, Concurrency: 10, Balance: 100},
	}
	return &geminiClientCancelFixture{handler: h, group: group, apiKey: apiKey, upstream: upstream, ctx: ctx}
}

// serve 经 ops 错误日志中间件执行请求，返回响应与 ops 队列中的条目数。
func (f *geminiClientCancelFixture) serve(t *testing.T, route, path, body string, call func(*gin.Context)) (*httptest.ResponseRecorder, int64) {
	t.Helper()
	setupOpsErrorLogTestQueue(t, 4)
	ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
	router := gin.New()
	router.Use(OpsErrorLoggerMiddleware(ops))
	router.POST(route, func(c *gin.Context) {
		c.Set(string(middleware.ContextKeyAPIKey), f.apiKey)
		c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: f.apiKey.UserID, Concurrency: 10})
		call(c)
	})

	req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(body))
	req = req.WithContext(context.WithValue(f.ctx, ctxkey.Group, f.group))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)
	return rec, OpsErrorLogQueueLength()
}

// 原生入口：上游响应前客户端断开，响应未提交，标记 499；纯客户端取消不落 ops 错误日志。
func TestGeminiV1BetaModels_ClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
	gin.SetMode(gin.TestMode)
	f := newGeminiClientCancelFixture(t)

	rec, opsQueued := f.serve(t,
		"/v1beta/models/*modelAction",
		"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
		`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
		f.handler.GeminiV1BetaModels,
	)

	require.Equal(t, int32(1), f.upstream.calls.Load())
	require.Equal(t, statusClientClosedRequest, rec.Code)
	require.Zero(t, rec.Body.Len())
	require.Zero(t, opsQueued)
}

// Chat Completions 兼容入口（Gemini 分组）：同一场景标记 499，不再补写 502。
func TestGatewayChatCompletions_GeminiClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
	gin.SetMode(gin.TestMode)
	f := newGeminiClientCancelFixture(t)

	rec, opsQueued := f.serve(t,
		"/v1/chat/completions",
		"/v1/chat/completions",
		`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}],"stream":true}`,
		f.handler.ChatCompletions,
	)

	require.Equal(t, int32(1), f.upstream.calls.Load())
	require.Equal(t, statusClientClosedRequest, rec.Code)
	require.Zero(t, rec.Body.Len())
	require.Zero(t, opsQueued)
}