195 lines
5.9 KiB
Go
195 lines
5.9 KiB
Go
|
//go:build unit
|
||
|
|
|
||
|
|
package handler
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
|
|
middleware2 "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"
|
||
|
|
"github.com/tidwall/gjson"
|
||
|
|
)
|
||
|
|
|
||
|
|
// openAIResponsesFailoverCancelUpstream 固定返回 HTTP 520，可在首次上游调用时
|
||
|
|
// 触发回调（用于模拟“上游在途期间客户端断开”）。
|
||
|
|
type openAIResponsesFailoverCancelUpstream struct {
|
||
|
|
service.HTTPUpstream
|
||
|
|
mu sync.Mutex
|
||
|
|
accountIDs []int64
|
||
|
|
onFirstDo func()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (u *openAIResponsesFailoverCancelUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||
|
|
u.mu.Lock()
|
||
|
|
u.accountIDs = append(u.accountIDs, accountID)
|
||
|
|
first := len(u.accountIDs) == 1
|
||
|
|
u.mu.Unlock()
|
||
|
|
if first && u.onFirstDo != nil {
|
||
|
|
u.onFirstDo()
|
||
|
|
}
|
||
|
|
return &http.Response{
|
||
|
|
StatusCode: 520,
|
||
|
|
Header: http.Header{"Content-Type": []string{"text/html"}},
|
||
|
|
Body: io.NopCloser(bytes.NewBufferString("<html>520: unknown error</html>")),
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (u *openAIResponsesFailoverCancelUpstream) calls() []int64 {
|
||
|
|
u.mu.Lock()
|
||
|
|
defer u.mu.Unlock()
|
||
|
|
return append([]int64(nil), u.accountIDs...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func newOpenAIResponsesFailoverTestHandler(t *testing.T, upstream service.HTTPUpstream) *OpenAIGatewayHandler {
|
||
|
|
t.Helper()
|
||
|
|
accounts := []service.Account{
|
||
|
|
{
|
||
|
|
ID: 1,
|
||
|
|
Name: "responses-account-1",
|
||
|
|
Platform: service.PlatformOpenAI,
|
||
|
|
Type: service.AccountTypeOAuth,
|
||
|
|
Status: service.StatusActive,
|
||
|
|
Schedulable: true,
|
||
|
|
Concurrency: 0,
|
||
|
|
Priority: 0,
|
||
|
|
Credentials: map[string]any{"access_token": "token-1"},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
ID: 2,
|
||
|
|
Name: "responses-account-2",
|
||
|
|
Platform: service.PlatformOpenAI,
|
||
|
|
Type: service.AccountTypeOAuth,
|
||
|
|
Status: service.StatusActive,
|
||
|
|
Schedulable: true,
|
||
|
|
Concurrency: 0,
|
||
|
|
Priority: 1,
|
||
|
|
Credentials: map[string]any{"access_token": "token-2"},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
accountRepo := openAIImagesFailoverAccountRepo{accounts: accounts}
|
||
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||
|
|
gatewayService := service.NewOpenAIGatewayService(
|
||
|
|
accountRepo,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
cfg,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
upstream,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
)
|
||
|
|
billingService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||
|
|
t.Cleanup(billingService.Stop)
|
||
|
|
concurrencyService := service.NewConcurrencyService(nil)
|
||
|
|
handler := NewOpenAIGatewayHandler(
|
||
|
|
gatewayService,
|
||
|
|
concurrencyService,
|
||
|
|
billingService,
|
||
|
|
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
nil,
|
||
|
|
cfg,
|
||
|
|
)
|
||
|
|
handler.maxAccountSwitches = 10
|
||
|
|
return handler
|
||
|
|
}
|
||
|
|
|
||
|
|
func newOpenAIResponsesFailoverTestContext(t *testing.T, ctx context.Context) (*gin.Context, *httptest.ResponseRecorder) {
|
||
|
|
t.Helper()
|
||
|
|
groupID := int64(3131)
|
||
|
|
body := []byte(`{"model":"gpt-5.1","stream":false,"input":"hello"}`)
|
||
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||
|
|
if ctx != nil {
|
||
|
|
req = req.WithContext(ctx)
|
||
|
|
}
|
||
|
|
req.Header.Set("Content-Type", "application/json")
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
c, _ := gin.CreateTestContext(rec)
|
||
|
|
c.Request = req
|
||
|
|
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||
|
|
ID: 99,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{
|
||
|
|
ID: groupID,
|
||
|
|
Platform: service.PlatformOpenAI,
|
||
|
|
},
|
||
|
|
User: &service.User{ID: 100},
|
||
|
|
})
|
||
|
|
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 100, Concurrency: 0})
|
||
|
|
return c, rec
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestOpenAIGatewayHandlerResponses_FailoverAbortsWhenClientDisconnected 复现
|
||
|
|
// #4257：客户端在上游请求在途期间断开，上游随后返回可 failover 的 520。
|
||
|
|
// 期望：不再用已取消的 context 重新选号（不触达账号 2）、不把取消误报成
|
||
|
|
// 502 账号耗尽、请求按 499 归类。
|
||
|
|
func TestOpenAIGatewayHandlerResponses_FailoverAbortsWhenClientDisconnected(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
defer cancel()
|
||
|
|
upstream := &openAIResponsesFailoverCancelUpstream{onFirstDo: cancel}
|
||
|
|
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||
|
|
c, rec := newOpenAIResponsesFailoverTestContext(t, ctx)
|
||
|
|
|
||
|
|
handler.Responses(c)
|
||
|
|
|
||
|
|
require.Equal(t, []int64{1}, upstream.calls(), "客户端断开后不应再切换到账号 2")
|
||
|
|
require.Equal(t, statusClientClosedRequest, c.Writer.Status(), "应按 499 归类")
|
||
|
|
require.Zero(t, rec.Body.Len(), "不应写入 502 错误响应体")
|
||
|
|
|
||
|
|
_, hasFinalUpstreamErr := c.Get(service.OpsUpstreamStatusCodeKey)
|
||
|
|
require.False(t, hasFinalUpstreamErr, "不应记录 failover 耗尽的上游错误终态")
|
||
|
|
|
||
|
|
// 真实发生过的 520 应保留 failover 事件（service 层在返回 failover 错误前记录）
|
||
|
|
rawEvents, ok := c.Get(service.OpsUpstreamErrorsKey)
|
||
|
|
require.True(t, ok)
|
||
|
|
events, ok := rawEvents.([]*service.OpsUpstreamErrorEvent)
|
||
|
|
require.True(t, ok)
|
||
|
|
require.Len(t, events, 1)
|
||
|
|
require.Equal(t, "failover", events[0].Kind)
|
||
|
|
require.Equal(t, 520, events[0].UpstreamStatusCode)
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestOpenAIGatewayHandlerResponses_FailoverContinuesForConnectedClient 回归
|
||
|
|
// 守卫：客户端在线时 failover 行为不变——切换到账号 2，两个账号都 520 后按
|
||
|
|
// 耗尽返回 502。
|
||
|
|
func TestOpenAIGatewayHandlerResponses_FailoverContinuesForConnectedClient(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
|
||
|
|
upstream := &openAIResponsesFailoverCancelUpstream{}
|
||
|
|
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||
|
|
c, rec := newOpenAIResponsesFailoverTestContext(t, nil)
|
||
|
|
|
||
|
|
handler.Responses(c)
|
||
|
|
|
||
|
|
require.Equal(t, []int64{1, 2}, upstream.calls(), "在线客户端应正常切换账号")
|
||
|
|
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||
|
|
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
|
||
|
|
}
|