33db7a0fb6
- 新增 UsageRecordWorkerPool，支持有界队列、溢出降级策略与自动扩缩容 - 将 Gateway/OpenAI/Sora/Gemini 使用量记录改为提交到统一任务池执行 - 增加 usage_record 配置默认值与校验规则，并补充配置与任务提交相关测试 - 注入并托管 worker 池生命周期，服务退出时统一 StopAndWait Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
689 lines
28 KiB
Go
689 lines
28 KiB
Go
//go:build unit
|
||
|
||
package handler
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
"github.com/Wei-Shaw/sub2api/internal/testutil"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/stretchr/testify/require"
|
||
"github.com/tidwall/gjson"
|
||
"github.com/tidwall/sjson"
|
||
)
|
||
|
||
// 编译期接口断言
|
||
var _ service.SoraClient = (*stubSoraClient)(nil)
|
||
var _ service.AccountRepository = (*stubAccountRepo)(nil)
|
||
var _ service.GroupRepository = (*stubGroupRepo)(nil)
|
||
var _ service.UsageLogRepository = (*stubUsageLogRepo)(nil)
|
||
|
||
type stubSoraClient struct {
|
||
imageURLs []string
|
||
}
|
||
|
||
func (s *stubSoraClient) Enabled() bool { return true }
|
||
func (s *stubSoraClient) UploadImage(ctx context.Context, account *service.Account, data []byte, filename string) (string, error) {
|
||
return "upload", nil
|
||
}
|
||
func (s *stubSoraClient) CreateImageTask(ctx context.Context, account *service.Account, req service.SoraImageRequest) (string, error) {
|
||
return "task-image", nil
|
||
}
|
||
func (s *stubSoraClient) CreateVideoTask(ctx context.Context, account *service.Account, req service.SoraVideoRequest) (string, error) {
|
||
return "task-video", nil
|
||
}
|
||
func (s *stubSoraClient) CreateStoryboardTask(ctx context.Context, account *service.Account, req service.SoraStoryboardRequest) (string, error) {
|
||
return "task-video", nil
|
||
}
|
||
func (s *stubSoraClient) UploadCharacterVideo(ctx context.Context, account *service.Account, data []byte) (string, error) {
|
||
return "cameo-1", nil
|
||
}
|
||
func (s *stubSoraClient) GetCameoStatus(ctx context.Context, account *service.Account, cameoID string) (*service.SoraCameoStatus, error) {
|
||
return &service.SoraCameoStatus{
|
||
Status: "finalized",
|
||
StatusMessage: "Completed",
|
||
DisplayNameHint: "Character",
|
||
UsernameHint: "user.character",
|
||
ProfileAssetURL: "https://example.com/avatar.webp",
|
||
}, nil
|
||
}
|
||
func (s *stubSoraClient) DownloadCharacterImage(ctx context.Context, account *service.Account, imageURL string) ([]byte, error) {
|
||
return []byte("avatar"), nil
|
||
}
|
||
func (s *stubSoraClient) UploadCharacterImage(ctx context.Context, account *service.Account, data []byte) (string, error) {
|
||
return "asset-pointer", nil
|
||
}
|
||
func (s *stubSoraClient) FinalizeCharacter(ctx context.Context, account *service.Account, req service.SoraCharacterFinalizeRequest) (string, error) {
|
||
return "character-1", nil
|
||
}
|
||
func (s *stubSoraClient) SetCharacterPublic(ctx context.Context, account *service.Account, cameoID string) error {
|
||
return nil
|
||
}
|
||
func (s *stubSoraClient) DeleteCharacter(ctx context.Context, account *service.Account, characterID string) error {
|
||
return nil
|
||
}
|
||
func (s *stubSoraClient) PostVideoForWatermarkFree(ctx context.Context, account *service.Account, generationID string) (string, error) {
|
||
return "s_post", nil
|
||
}
|
||
func (s *stubSoraClient) DeletePost(ctx context.Context, account *service.Account, postID string) error {
|
||
return nil
|
||
}
|
||
func (s *stubSoraClient) GetWatermarkFreeURLCustom(ctx context.Context, account *service.Account, parseURL, parseToken, postID string) (string, error) {
|
||
return "https://example.com/no-watermark.mp4", nil
|
||
}
|
||
func (s *stubSoraClient) EnhancePrompt(ctx context.Context, account *service.Account, prompt, expansionLevel string, durationS int) (string, error) {
|
||
return "enhanced prompt", nil
|
||
}
|
||
func (s *stubSoraClient) GetImageTask(ctx context.Context, account *service.Account, taskID string) (*service.SoraImageTaskStatus, error) {
|
||
return &service.SoraImageTaskStatus{ID: taskID, Status: "completed", URLs: s.imageURLs}, nil
|
||
}
|
||
func (s *stubSoraClient) GetVideoTask(ctx context.Context, account *service.Account, taskID string) (*service.SoraVideoTaskStatus, error) {
|
||
return &service.SoraVideoTaskStatus{ID: taskID, Status: "completed", URLs: s.imageURLs}, nil
|
||
}
|
||
|
||
type stubAccountRepo struct {
|
||
accounts map[int64]*service.Account
|
||
}
|
||
|
||
func (r *stubAccountRepo) Create(ctx context.Context, account *service.Account) error { return nil }
|
||
func (r *stubAccountRepo) GetByID(ctx context.Context, id int64) (*service.Account, error) {
|
||
if acc, ok := r.accounts[id]; ok {
|
||
return acc, nil
|
||
}
|
||
return nil, service.ErrAccountNotFound
|
||
}
|
||
func (r *stubAccountRepo) GetByIDs(ctx context.Context, ids []int64) ([]*service.Account, error) {
|
||
var result []*service.Account
|
||
for _, id := range ids {
|
||
if acc, ok := r.accounts[id]; ok {
|
||
result = append(result, acc)
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
func (r *stubAccountRepo) ExistsByID(ctx context.Context, id int64) (bool, error) {
|
||
_, ok := r.accounts[id]
|
||
return ok, nil
|
||
}
|
||
func (r *stubAccountRepo) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*service.Account, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubAccountRepo) FindByExtraField(ctx context.Context, key string, value any) ([]service.Account, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubAccountRepo) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
|
||
return map[string]int64{}, nil
|
||
}
|
||
func (r *stubAccountRepo) Update(ctx context.Context, account *service.Account) error { return nil }
|
||
func (r *stubAccountRepo) Delete(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.Account, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (r *stubAccountRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]service.Account, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (r *stubAccountRepo) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubAccountRepo) ListActive(ctx context.Context) ([]service.Account, error) { return nil, nil }
|
||
func (r *stubAccountRepo) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
|
||
return r.listSchedulableByPlatform(platform), nil
|
||
}
|
||
func (r *stubAccountRepo) UpdateLastUsed(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) SetError(ctx context.Context, id int64, errorMsg string) error { return nil }
|
||
func (r *stubAccountRepo) ClearError(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
|
||
return 0, nil
|
||
}
|
||
func (r *stubAccountRepo) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulable(ctx context.Context) ([]service.Account, error) {
|
||
return r.listSchedulable(), nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||
return r.listSchedulable(), nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
|
||
return r.listSchedulableByPlatform(platform), nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
|
||
return r.listSchedulableByPlatform(platform), nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
|
||
var result []service.Account
|
||
for _, acc := range r.accounts {
|
||
for _, platform := range platforms {
|
||
if acc.Platform == platform && acc.IsSchedulable() {
|
||
result = append(result, *acc)
|
||
break
|
||
}
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
func (r *stubAccountRepo) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
|
||
return r.ListSchedulableByPlatforms(ctx, platforms)
|
||
}
|
||
func (r *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) ClearTempUnschedulable(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) ClearRateLimit(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) ClearModelRateLimits(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubAccountRepo) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||
return nil
|
||
}
|
||
func (r *stubAccountRepo) BulkUpdate(ctx context.Context, ids []int64, updates service.AccountBulkUpdate) (int64, error) {
|
||
return 0, nil
|
||
}
|
||
|
||
func (r *stubAccountRepo) listSchedulable() []service.Account {
|
||
var result []service.Account
|
||
for _, acc := range r.accounts {
|
||
if acc.IsSchedulable() {
|
||
result = append(result, *acc)
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
func (r *stubAccountRepo) listSchedulableByPlatform(platform string) []service.Account {
|
||
var result []service.Account
|
||
for _, acc := range r.accounts {
|
||
if acc.Platform == platform && acc.IsSchedulable() {
|
||
result = append(result, *acc)
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
type stubGroupRepo struct {
|
||
group *service.Group
|
||
}
|
||
|
||
func (r *stubGroupRepo) Create(ctx context.Context, group *service.Group) error { return nil }
|
||
func (r *stubGroupRepo) GetByID(ctx context.Context, id int64) (*service.Group, error) {
|
||
return r.group, nil
|
||
}
|
||
func (r *stubGroupRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||
return r.group, nil
|
||
}
|
||
func (r *stubGroupRepo) Update(ctx context.Context, group *service.Group) error { return nil }
|
||
func (r *stubGroupRepo) Delete(ctx context.Context, id int64) error { return nil }
|
||
func (r *stubGroupRepo) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubGroupRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.Group, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (r *stubGroupRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (r *stubGroupRepo) ListActive(ctx context.Context) ([]service.Group, error) { return nil, nil }
|
||
func (r *stubGroupRepo) ListActiveByPlatform(ctx context.Context, platform string) ([]service.Group, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubGroupRepo) ExistsByName(ctx context.Context, name string) (bool, error) {
|
||
return false, nil
|
||
}
|
||
func (r *stubGroupRepo) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||
return 0, nil
|
||
}
|
||
func (r *stubGroupRepo) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||
return 0, nil
|
||
}
|
||
func (r *stubGroupRepo) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error) {
|
||
return nil, nil
|
||
}
|
||
func (r *stubGroupRepo) BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error {
|
||
return nil
|
||
}
|
||
func (r *stubGroupRepo) UpdateSortOrders(ctx context.Context, updates []service.GroupSortOrderUpdate) error {
|
||
return nil
|
||
}
|
||
|
||
type stubUsageLogRepo struct{}
|
||
|
||
func (s *stubUsageLogRepo) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
|
||
return true, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetByID(ctx context.Context, id int64) (*service.UsageLog, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) Delete(ctx context.Context, id int64) error { return nil }
|
||
func (s *stubUsageLogRepo) ListByUser(ctx context.Context, userID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByAPIKey(ctx context.Context, apiKeyID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByAccount(ctx context.Context, accountID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByUserAndTimeRange(ctx context.Context, userID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByAPIKeyAndTimeRange(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByAccountAndTimeRange(ctx context.Context, accountID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListByModelAndTimeRange(ctx context.Context, modelName string, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAccountTodayStats(ctx context.Context, accountID int64) (*usagestats.AccountStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetDashboardStats(ctx context.Context) (*usagestats.DashboardStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, stream *bool, billingType *int8) ([]usagestats.TrendDataPoint, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, stream *bool, billingType *int8) ([]usagestats.ModelStat, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.APIKeyUsageTrendPoint, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetBatchUserUsageStats(ctx context.Context, userIDs []int64, startTime, endTime time.Time) (map[int64]*usagestats.BatchUserUsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetBatchAPIKeyUsageStats(ctx context.Context, apiKeyIDs []int64, startTime, endTime time.Time) (map[int64]*usagestats.BatchAPIKeyUsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUserDashboardStats(ctx context.Context, userID int64) (*usagestats.UserDashboardStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID int64) (*usagestats.UserDashboardStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) ([]usagestats.TrendDataPoint, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) ([]usagestats.ModelStat, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters usagestats.UsageLogFilters) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||
return nil, nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetStatsWithFilters(ctx context.Context, filters usagestats.UsageLogFilters) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAccountUsageStats(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.AccountUsageStatsResponse, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAPIKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetAccountStatsAggregated(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetModelStatsAggregated(ctx context.Context, modelName string, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||
return nil, nil
|
||
}
|
||
func (s *stubUsageLogRepo) GetDailyStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) ([]map[string]any, error) {
|
||
return nil, nil
|
||
}
|
||
|
||
func TestSoraGatewayHandler_ChatCompletions(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
RunMode: config.RunModeSimple,
|
||
Gateway: config.GatewayConfig{
|
||
SoraStreamMode: "force",
|
||
MaxAccountSwitches: 1,
|
||
Scheduling: config.GatewaySchedulingConfig{
|
||
LoadBatchEnabled: false,
|
||
},
|
||
},
|
||
Concurrency: config.ConcurrencyConfig{PingInterval: 0},
|
||
Sora: config.SoraConfig{
|
||
Client: config.SoraClientConfig{
|
||
BaseURL: "https://sora.test",
|
||
PollIntervalSeconds: 1,
|
||
MaxPollAttempts: 1,
|
||
},
|
||
},
|
||
}
|
||
|
||
account := &service.Account{ID: 1, Platform: service.PlatformSora, Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 1}
|
||
accountRepo := &stubAccountRepo{accounts: map[int64]*service.Account{account.ID: account}}
|
||
group := &service.Group{ID: 1, Platform: service.PlatformSora, Status: service.StatusActive, Hydrated: true}
|
||
groupRepo := &stubGroupRepo{group: group}
|
||
|
||
usageLogRepo := &stubUsageLogRepo{}
|
||
deferredService := service.NewDeferredService(accountRepo, nil, 0)
|
||
billingService := service.NewBillingService(cfg, nil)
|
||
concurrencyService := service.NewConcurrencyService(testutil.StubConcurrencyCache{})
|
||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, cfg)
|
||
t.Cleanup(func() {
|
||
billingCacheService.Stop()
|
||
})
|
||
|
||
gatewayService := service.NewGatewayService(
|
||
accountRepo,
|
||
groupRepo,
|
||
usageLogRepo,
|
||
nil,
|
||
nil,
|
||
nil,
|
||
testutil.StubGatewayCache{},
|
||
cfg,
|
||
nil,
|
||
concurrencyService,
|
||
billingService,
|
||
nil,
|
||
billingCacheService,
|
||
nil,
|
||
nil,
|
||
deferredService,
|
||
nil,
|
||
testutil.StubSessionLimitCache{},
|
||
nil,
|
||
)
|
||
|
||
soraClient := &stubSoraClient{imageURLs: []string{"https://example.com/a.png"}}
|
||
soraGatewayService := service.NewSoraGatewayService(soraClient, nil, nil, cfg)
|
||
|
||
handler := NewSoraGatewayHandler(gatewayService, soraGatewayService, concurrencyService, billingCacheService, nil, cfg)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
body := `{"model":"gpt-image","messages":[{"role":"user","content":"hello"}]}`
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/sora/v1/chat/completions", strings.NewReader(body))
|
||
c.Request.Header.Set("Content-Type", "application/json")
|
||
|
||
apiKey := &service.APIKey{
|
||
ID: 1,
|
||
UserID: 1,
|
||
Status: service.StatusActive,
|
||
GroupID: &group.ID,
|
||
User: &service.User{ID: 1, Concurrency: 1, Status: service.StatusActive},
|
||
Group: group,
|
||
}
|
||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: apiKey.User.Concurrency})
|
||
|
||
handler.ChatCompletions(c)
|
||
|
||
require.Equal(t, http.StatusOK, rec.Code)
|
||
var resp map[string]any
|
||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||
require.NotEmpty(t, resp["media_url"])
|
||
}
|
||
|
||
// TestSoraHandler_StreamForcing 验证 sora handler 的 stream 强制逻辑
|
||
func TestSoraHandler_StreamForcing(t *testing.T) {
|
||
// 测试 1：stream=false 时 sjson 强制修改为 true
|
||
body := []byte(`{"model":"sora","messages":[{"role":"user","content":"test"}],"stream":false}`)
|
||
clientStream := gjson.GetBytes(body, "stream").Bool()
|
||
require.False(t, clientStream)
|
||
newBody, err := sjson.SetBytes(body, "stream", true)
|
||
require.NoError(t, err)
|
||
require.True(t, gjson.GetBytes(newBody, "stream").Bool())
|
||
|
||
// 测试 2：stream=true 时不修改
|
||
body2 := []byte(`{"model":"sora","messages":[{"role":"user","content":"test"}],"stream":true}`)
|
||
require.True(t, gjson.GetBytes(body2, "stream").Bool())
|
||
|
||
// 测试 3：无 stream 字段时 gjson 返回 false（零值）
|
||
body3 := []byte(`{"model":"sora","messages":[{"role":"user","content":"test"}]}`)
|
||
require.False(t, gjson.GetBytes(body3, "stream").Bool())
|
||
}
|
||
|
||
// TestSoraHandler_ValidationExtraction 验证 sora handler 中 gjson 字段校验逻辑
|
||
func TestSoraHandler_ValidationExtraction(t *testing.T) {
|
||
// model 缺失
|
||
body := []byte(`{"messages":[{"role":"user","content":"test"}]}`)
|
||
modelResult := gjson.GetBytes(body, "model")
|
||
require.True(t, !modelResult.Exists() || modelResult.Type != gjson.String || modelResult.String() == "")
|
||
|
||
// model 为数字 → 类型不是 gjson.String，应被拒绝
|
||
body1b := []byte(`{"model":123,"messages":[{"role":"user","content":"test"}]}`)
|
||
modelResult1b := gjson.GetBytes(body1b, "model")
|
||
require.True(t, modelResult1b.Exists())
|
||
require.NotEqual(t, gjson.String, modelResult1b.Type)
|
||
|
||
// messages 缺失
|
||
body2 := []byte(`{"model":"sora"}`)
|
||
require.False(t, gjson.GetBytes(body2, "messages").IsArray())
|
||
|
||
// messages 不是 JSON 数组（字符串）
|
||
body3 := []byte(`{"model":"sora","messages":"not array"}`)
|
||
require.False(t, gjson.GetBytes(body3, "messages").IsArray())
|
||
|
||
// messages 是对象而非数组 → IsArray 返回 false
|
||
body4 := []byte(`{"model":"sora","messages":{}}`)
|
||
require.False(t, gjson.GetBytes(body4, "messages").IsArray())
|
||
|
||
// messages 是空数组 → IsArray 为 true 但 len==0，应被拒绝
|
||
body5 := []byte(`{"model":"sora","messages":[]}`)
|
||
msgsResult := gjson.GetBytes(body5, "messages")
|
||
require.True(t, msgsResult.IsArray())
|
||
require.Equal(t, 0, len(msgsResult.Array()))
|
||
|
||
// 非法 JSON 被 gjson.ValidBytes 拦截
|
||
require.False(t, gjson.ValidBytes([]byte(`{invalid`)))
|
||
}
|
||
|
||
// TestGenerateOpenAISessionHash_WithBody 验证 generateOpenAISessionHash 的 body/header 解析逻辑
|
||
func TestGenerateOpenAISessionHash_WithBody(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
// 从 body 提取 prompt_cache_key
|
||
body := []byte(`{"model":"sora","prompt_cache_key":"session-abc"}`)
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
c.Request = httptest.NewRequest("POST", "/", nil)
|
||
|
||
hash := generateOpenAISessionHash(c, body)
|
||
require.NotEmpty(t, hash)
|
||
|
||
// 无 prompt_cache_key 且无 header → 空 hash
|
||
body2 := []byte(`{"model":"sora"}`)
|
||
hash2 := generateOpenAISessionHash(c, body2)
|
||
require.Empty(t, hash2)
|
||
|
||
// header 优先于 body
|
||
c.Request.Header.Set("session_id", "from-header")
|
||
hash3 := generateOpenAISessionHash(c, body)
|
||
require.NotEmpty(t, hash3)
|
||
require.NotEqual(t, hash, hash3) // 不同来源应产生不同 hash
|
||
}
|
||
|
||
func TestSoraHandleStreamingAwareError_JSONEscaping(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
errType string
|
||
message string
|
||
}{
|
||
{
|
||
name: "包含双引号",
|
||
errType: "upstream_error",
|
||
message: `upstream returned "invalid" payload`,
|
||
},
|
||
{
|
||
name: "包含换行和制表符",
|
||
errType: "rate_limit_error",
|
||
message: "line1\nline2\ttab",
|
||
},
|
||
{
|
||
name: "包含反斜杠",
|
||
errType: "upstream_error",
|
||
message: `path C:\Users\test\file.txt not found`,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
||
h := &SoraGatewayHandler{}
|
||
h.handleStreamingAwareError(c, http.StatusBadGateway, tt.errType, tt.message, true)
|
||
|
||
body := w.Body.String()
|
||
require.True(t, strings.HasPrefix(body, "event: error\n"), "应以 SSE error 事件开头")
|
||
require.True(t, strings.HasSuffix(body, "\n\n"), "应以 SSE 结束分隔符结尾")
|
||
|
||
lines := strings.Split(strings.TrimSuffix(body, "\n\n"), "\n")
|
||
require.Len(t, lines, 2, "SSE 错误事件应包含 event 行和 data 行")
|
||
require.Equal(t, "event: error", lines[0])
|
||
require.True(t, strings.HasPrefix(lines[1], "data: "), "第二行应为 data 前缀")
|
||
|
||
jsonStr := strings.TrimPrefix(lines[1], "data: ")
|
||
var parsed map[string]any
|
||
require.NoError(t, json.Unmarshal([]byte(jsonStr), &parsed), "data 行必须是合法 JSON")
|
||
|
||
errorObj, ok := parsed["error"].(map[string]any)
|
||
require.True(t, ok, "JSON 中应包含 error 对象")
|
||
require.Equal(t, tt.errType, errorObj["type"])
|
||
require.Equal(t, tt.message, errorObj["message"])
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestSoraHandleFailoverExhausted_StreamPassesUpstreamMessage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
||
h := &SoraGatewayHandler{}
|
||
resp := []byte(`{"error":{"message":"invalid \"prompt\"\nline2","code":"bad_request"}}`)
|
||
h.handleFailoverExhausted(c, http.StatusBadGateway, nil, resp, true)
|
||
|
||
body := w.Body.String()
|
||
require.True(t, strings.HasPrefix(body, "event: error\n"))
|
||
require.True(t, strings.HasSuffix(body, "\n\n"))
|
||
|
||
lines := strings.Split(strings.TrimSuffix(body, "\n\n"), "\n")
|
||
require.Len(t, lines, 2)
|
||
jsonStr := strings.TrimPrefix(lines[1], "data: ")
|
||
|
||
var parsed map[string]any
|
||
require.NoError(t, json.Unmarshal([]byte(jsonStr), &parsed))
|
||
|
||
errorObj, ok := parsed["error"].(map[string]any)
|
||
require.True(t, ok)
|
||
require.Equal(t, "upstream_error", errorObj["type"])
|
||
require.Equal(t, "invalid \"prompt\"\nline2", errorObj["message"])
|
||
}
|
||
|
||
func TestSoraHandleFailoverExhausted_CloudflareChallengeIncludesRay(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
||
headers := http.Header{}
|
||
headers.Set("cf-ray", "9d01b0e9ecc35829-SEA")
|
||
body := []byte(`<!DOCTYPE html><html><head><title>Just a moment...</title></head><body><script>window._cf_chl_opt={};</script></body></html>`)
|
||
|
||
h := &SoraGatewayHandler{}
|
||
h.handleFailoverExhausted(c, http.StatusForbidden, headers, body, true)
|
||
|
||
lines := strings.Split(strings.TrimSuffix(w.Body.String(), "\n\n"), "\n")
|
||
require.Len(t, lines, 2)
|
||
jsonStr := strings.TrimPrefix(lines[1], "data: ")
|
||
|
||
var parsed map[string]any
|
||
require.NoError(t, json.Unmarshal([]byte(jsonStr), &parsed))
|
||
|
||
errorObj, ok := parsed["error"].(map[string]any)
|
||
require.True(t, ok)
|
||
require.Equal(t, "upstream_error", errorObj["type"])
|
||
msg, _ := errorObj["message"].(string)
|
||
require.Contains(t, msg, "Cloudflare challenge")
|
||
require.Contains(t, msg, "cf-ray: 9d01b0e9ecc35829-SEA")
|
||
}
|
||
|
||
func TestSoraHandleFailoverExhausted_CfShield429MappedToRateLimitError(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
||
headers := http.Header{}
|
||
headers.Set("cf-ray", "9d03b68c086027a1-SEA")
|
||
body := []byte(`{"error":{"code":"cf_shield_429","message":"shield blocked"}}`)
|
||
|
||
h := &SoraGatewayHandler{}
|
||
h.handleFailoverExhausted(c, http.StatusTooManyRequests, headers, body, true)
|
||
|
||
lines := strings.Split(strings.TrimSuffix(w.Body.String(), "\n\n"), "\n")
|
||
require.Len(t, lines, 2)
|
||
jsonStr := strings.TrimPrefix(lines[1], "data: ")
|
||
|
||
var parsed map[string]any
|
||
require.NoError(t, json.Unmarshal([]byte(jsonStr), &parsed))
|
||
|
||
errorObj, ok := parsed["error"].(map[string]any)
|
||
require.True(t, ok)
|
||
require.Equal(t, "rate_limit_error", errorObj["type"])
|
||
msg, _ := errorObj["message"].(string)
|
||
require.Contains(t, msg, "Cloudflare shield")
|
||
require.Contains(t, msg, "cf-ray: 9d03b68c086027a1-SEA")
|
||
}
|
||
|
||
func TestExtractSoraFailoverHeaderInsights(t *testing.T) {
|
||
headers := http.Header{}
|
||
headers.Set("cf-mitigated", "challenge")
|
||
headers.Set("content-type", "text/html")
|
||
body := []byte(`<script>window._cf_chl_opt={cRay: '9cff2d62d83bb98d'};</script>`)
|
||
|
||
rayID, mitigated, contentType := extractSoraFailoverHeaderInsights(headers, body)
|
||
require.Equal(t, "9cff2d62d83bb98d", rayID)
|
||
require.Equal(t, "challenge", mitigated)
|
||
require.Equal(t, "text/html", contentType)
|
||
}
|