Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/account_scheduling_threshold_integration_test.go

@@ -0,0 +1,136 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/stretchr/testify/require"
+)
+
+type thresholdSelectionAccountRepoStub struct {
+	rateLimitAccountRepoStub
+	accounts []Account
+}
+
+func (r *thresholdSelectionAccountRepoStub) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) {
+	filtered := make([]Account, 0, len(r.accounts))
+	for _, account := range r.accounts {
+		if account.Platform == platform {
+			filtered = append(filtered, account)
+		}
+	}
+	return filtered, nil
+}
+
+func (r *thresholdSelectionAccountRepoStub) ListSchedulableByGroupIDAndPlatform(ctx context.Context, _ int64, platform string) ([]Account, error) {
+	return r.ListSchedulableByPlatform(ctx, platform)
+}
+
+func (r *thresholdSelectionAccountRepoStub) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
+	return r.ListSchedulableByPlatform(ctx, platform)
+}
+
+func TestGatewayService_ListSchedulableAccounts_DoesNotFilterUnsupportedThresholdPlatforms(t *testing.T) {
+	accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
+	accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
+
+	settingsRepo := newMockSettingRepo()
+	settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":90}`
+
+	accountRepo := &thresholdSelectionAccountRepoStub{
+		accounts: []Account{
+			{
+				ID:          3101,
+				Platform:    PlatformKiro,
+				Status:      StatusActive,
+				Schedulable: true,
+				Credentials: map[string]any{
+					"account_scheduling_threshold": 1,
+				},
+				Extra: map[string]any{
+					"kiro_sched_utilization": 95.0,
+					"kiro_sched_reset_at":    time.Now().UTC().Add(2 * time.Hour).Format(time.RFC3339),
+				},
+			},
+			{
+				ID:          3102,
+				Platform:    PlatformKiro,
+				Status:      StatusActive,
+				Schedulable: true,
+				Extra: map[string]any{
+					"kiro_sched_utilization": 42.0,
+					"kiro_sched_reset_at":    time.Now().UTC().Add(2 * time.Hour).Format(time.RFC3339),
+				},
+			},
+		},
+	}
+
+	rateLimitService := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
+	rateLimitService.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
+	svc := &GatewayService{
+		accountRepo:      accountRepo,
+		cfg:              &config.Config{},
+		rateLimitService: rateLimitService,
+	}
+
+	accounts, useMixed, err := svc.listSchedulableAccounts(context.Background(), nil, PlatformKiro, false)
+
+	require.NoError(t, err)
+	require.False(t, useMixed)
+	require.Len(t, accounts, 2)
+	require.Equal(t, int64(3101), accounts[0].ID)
+	require.Equal(t, int64(3102), accounts[1].ID)
+	require.Equal(t, 0, accountRepo.tempCalls)
+}
+
+func TestOpenAIGatewayService_ListSchedulableAccounts_FiltersThresholdBlockedAccounts(t *testing.T) {
+	accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
+	accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
+
+	settingsRepo := newMockSettingRepo()
+	settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":85}`
+
+	accountRepo := &thresholdSelectionAccountRepoStub{
+		accounts: []Account{
+			{
+				ID:          4101,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Extra: map[string]any{
+					"codex_7d_used_percent": 91.0,
+					"codex_7d_reset_at":     time.Now().UTC().Add(12 * time.Hour).Format(time.RFC3339),
+				},
+			},
+			{
+				ID:          4102,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Extra: map[string]any{
+					"codex_7d_used_percent": 40.0,
+					"codex_7d_reset_at":     time.Now().UTC().Add(12 * time.Hour).Format(time.RFC3339),
+				},
+			},
+		},
+	}
+
+	rateLimitService := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
+	rateLimitService.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
+	svc := &OpenAIGatewayService{
+		accountRepo:      accountRepo,
+		cfg:              &config.Config{},
+		rateLimitService: rateLimitService,
+	}
+
+	accounts, err := svc.listSchedulableAccounts(context.Background(), nil, PlatformOpenAI)
+
+	require.NoError(t, err)
+	require.Len(t, accounts, 1)
+	require.Equal(t, int64(4102), accounts[0].ID)
+	require.Equal(t, 1, accountRepo.tempCalls)
+}