Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s

This commit is contained in:
李建琦
2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,118 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAccount_IsAnthropicAPIKeyPassthroughEnabled(t *testing.T) {
t.Run("Anthropic API Key 开启", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_passthrough": true,
},
}
require.True(t, account.IsAnthropicAPIKeyPassthroughEnabled())
})
t.Run("Anthropic API Key 关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_passthrough": false,
},
}
require.False(t, account.IsAnthropicAPIKeyPassthroughEnabled())
})
t.Run("字段类型非法默认关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_passthrough": "true",
},
}
require.False(t, account.IsAnthropicAPIKeyPassthroughEnabled())
})
t.Run("非 Anthropic API Key 账号始终关闭", func(t *testing.T) {
oauth := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"anthropic_passthrough": true,
},
}
require.False(t, oauth.IsAnthropicAPIKeyPassthroughEnabled())
openai := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_passthrough": true,
},
}
require.False(t, openai.IsAnthropicAPIKeyPassthroughEnabled())
})
}
func TestAccount_GetAnthropicAPIKeyAuthScheme(t *testing.T) {
tests := []struct {
name string
account *Account
want string
}{
{
name: "missing extra defaults to x-api-key",
account: &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
},
want: AnthropicAPIKeyAuthSchemeXAPIKey,
},
{
name: "explicit bearer",
account: &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_apikey_auth_scheme": AnthropicAPIKeyAuthSchemeAuthorizationBearer,
},
},
want: AnthropicAPIKeyAuthSchemeAuthorizationBearer,
},
{
name: "invalid value defaults to x-api-key",
account: &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_apikey_auth_scheme": "bearer",
},
},
want: AnthropicAPIKeyAuthSchemeXAPIKey,
},
{
name: "non Anthropic API key defaults to x-api-key",
account: &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"anthropic_apikey_auth_scheme": AnthropicAPIKeyAuthSchemeAuthorizationBearer,
},
},
want: AnthropicAPIKeyAuthSchemeXAPIKey,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, tt.account.GetAnthropicAPIKeyAuthScheme())
})
}
}
@@ -0,0 +1,398 @@
//go:build unit
package service
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
func TestGetBaseURL(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "non-apikey type returns empty",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAnthropic,
},
expected: "",
},
{
name: "apikey without base_url returns default anthropic",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAnthropic,
Credentials: map[string]any{},
},
expected: "https://api.anthropic.com",
},
{
name: "apikey with custom base_url",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAnthropic,
Credentials: map[string]any{"base_url": "https://custom.example.com"},
},
expected: "https://custom.example.com",
},
{
name: "antigravity apikey auto-appends /antigravity",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity apikey trims trailing slash before appending",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity non-apikey returns empty",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetBaseURL()
if result != tt.expected {
t.Errorf("GetBaseURL() = %q, want %q", result, tt.expected)
}
})
}
}
func TestGetGeminiBaseURL(t *testing.T) {
const defaultGeminiURL = "https://generativelanguage.googleapis.com"
tests := []struct {
name string
account Account
expected string
}{
{
name: "apikey without base_url returns default",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{},
},
expected: defaultGeminiURL,
},
{
name: "apikey with custom base_url",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{"base_url": "https://custom-gemini.example.com"},
},
expected: "https://custom-gemini.example.com",
},
{
name: "antigravity apikey auto-appends /antigravity",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity apikey trims trailing slash",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity oauth does NOT append /antigravity",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com",
},
{
name: "oauth without base_url returns default",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{},
},
expected: defaultGeminiURL,
},
{
name: "nil credentials returns default",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
},
expected: defaultGeminiURL,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetGeminiBaseURL(defaultGeminiURL)
if result != tt.expected {
t.Errorf("GetGeminiBaseURL() = %q, want %q", result, tt.expected)
}
})
}
}
func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "oauth without base_url uses CLI subscription proxy",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth stored official API endpoint is honored (manual endpoint switch)",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL,
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored regional API endpoint is honored",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://us-west-2.api.x.ai/v1",
},
},
expected: "https://us-west-2.api.x.ai/v1",
},
{
name: "oauth stored CLI proxy is honored verbatim",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultCLIBaseURL,
},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth unparseable base_url falls back to CLI proxy",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "not a url",
},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth explicit custom base_url redirects forwarding traffic",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
},
expected: "https://custom.example.com/v1",
},
{
name: "oauth custom base_url with path prefix redirects forwarding traffic",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://relay.example.com/xai/v1",
},
},
expected: "https://relay.example.com/xai/v1",
},
{
name: "API key without base_url uses official credit-backed API",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultBaseURL,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetGrokBaseURL())
})
}
}
func TestGetGrokBaseURLHonorsOAuthCustomRegardlessOfUnsafeOverrides(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
account := Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
}
require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL())
}
func TestGetGrokMediaBaseURLRedirectsCLIGatewayToOfficialAPI(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "oauth without base_url uses official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored CLI proxy is separated from the media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultCLIBaseURL,
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored CLI proxy variant is canonicalized to the media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/",
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth unparseable base_url falls back to official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "not a url",
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored official API endpoint is honored (manual endpoint switch)",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL,
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored regional API endpoint is honored for media",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://us-west-2.api.x.ai/v1",
},
},
expected: "https://us-west-2.api.x.ai/v1",
},
{
name: "oauth custom base_url redirects media traffic",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
},
expected: "https://custom.example.com/v1",
},
{
name: "API key retains its configured media API",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://grok.example.com/v1",
},
},
expected: "https://grok.example.com/v1",
},
{
name: "non-Grok account has no Grok media base URL",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformOpenAI,
Credentials: map[string]any{},
},
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetGrokMediaBaseURL())
})
}
}
func TestGetGrokMediaBaseURLHonorsOAuthCustomRegardlessOfUnsafeOverrides(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
account := Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
}
require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL())
}
@@ -0,0 +1,27 @@
package service
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestAccount_BillingRateMultiplier_DefaultsToOneWhenNil(t *testing.T) {
var a Account
require.NoError(t, json.Unmarshal([]byte(`{"id":1,"name":"acc","status":"active"}`), &a))
require.Nil(t, a.RateMultiplier)
require.Equal(t, 1.0, a.BillingRateMultiplier())
}
func TestAccount_BillingRateMultiplier_AllowsZero(t *testing.T) {
v := 0.0
a := Account{RateMultiplier: &v}
require.Equal(t, 0.0, a.BillingRateMultiplier())
}
func TestAccount_BillingRateMultiplier_NegativeFallsBackToOne(t *testing.T) {
v := -1.0
a := Account{RateMultiplier: &v}
require.Equal(t, 1.0, a.BillingRateMultiplier())
}
@@ -0,0 +1,45 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAccount_IsCodexCLIOnlyAppServerAllowed(t *testing.T) {
t.Run("codex_cli_only 开 + allow_app_server=true → true", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{"codex_cli_only": true, "codex_cli_only_allow_app_server": true},
}
require.True(t, account.IsCodexCLIOnlyAppServerAllowed())
})
t.Run("codex_cli_only 开 + allow_app_server=false → false", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{"codex_cli_only": true, "codex_cli_only_allow_app_server": false},
}
require.False(t, account.IsCodexCLIOnlyAppServerAllowed())
})
t.Run("codex_cli_only 开 + 字段缺失 → false", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{"codex_cli_only": true},
}
require.False(t, account.IsCodexCLIOnlyAppServerAllowed())
})
t.Run("codex_cli_only 关 → 即便 allow_app_server=true 也 false", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{"codex_cli_only_allow_app_server": true},
}
require.False(t, account.IsCodexCLIOnlyAppServerAllowed())
})
}
@@ -0,0 +1,245 @@
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
const userSuppliedCodexFingerprintSeed = "22222222-2222-4222-8222-222222222222"
func requireValidCodexFingerprintSeed(t *testing.T, extra map[string]any) string {
t.Helper()
seed, ok := codexFingerprintSeed(extra)
require.True(t, ok, "expected valid canonical Codex fingerprint seed")
return seed
}
func TestAdminCreateAccountStripsUserSeedAndCreatesFreshSeedWhenEnabled(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
svc := &adminServiceImpl{accountRepo: repo}
created, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
Name: "codex-oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
SkipDefaultGroupBind: true,
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: userSuppliedCodexFingerprintSeed,
},
})
require.NoError(t, err)
seed := requireValidCodexFingerprintSeed(t, created.Extra)
require.NotEqual(t, userSuppliedCodexFingerprintSeed, seed)
require.Equal(t, "session", created.Extra[codexFingerprintModeExtraKey])
}
func TestAdminUpdateAccountPreservesExistingSeedAndStripsUserSeed(t *testing.T) {
accountID := int64(201)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Name: "before",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: testCodexFingerprintSeed,
},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
codexFingerprintModeExtraKey: "full",
codexFingerprintSeedExtraKey: userSuppliedCodexFingerprintSeed,
"custom": "value",
},
})
require.NoError(t, err)
require.Equal(t, testCodexFingerprintSeed, requireValidCodexFingerprintSeed(t, updated.Extra))
require.Equal(t, "full", updated.Extra[codexFingerprintModeExtraKey])
require.Equal(t, "value", updated.Extra["custom"])
}
func TestAdminUpdateAccountInitializesSeedWhenFullEditEnables(t *testing.T) {
accountID := int64(202)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Name: "before",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{
codexFingerprintModeExtraKey: "off",
codexFingerprintSeedExtraKey: "not-a-seed",
},
},
}}
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{codexFingerprintModeExtraKey: "device"},
})
require.NoError(t, err)
require.NotEqual(t, "not-a-seed", requireValidCodexFingerprintSeed(t, updated.Extra))
require.Equal(t, "device", updated.Extra[codexFingerprintModeExtraKey])
}
func TestAdminUpdateAccountDisableReenablePreservesValidSeed(t *testing.T) {
accountID := int64(203)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: testCodexFingerprintSeed,
},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
disabled, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{codexFingerprintModeExtraKey: "off"},
})
require.NoError(t, err)
require.Equal(t, testCodexFingerprintSeed, requireValidCodexFingerprintSeed(t, disabled.Extra))
reenabled, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{codexFingerprintModeExtraKey: "session"},
})
require.NoError(t, err)
require.Equal(t, testCodexFingerprintSeed, requireValidCodexFingerprintSeed(t, reenabled.Extra))
}
func TestAdminUpdateAccountExtraStripsSeedAndLeavesAtomicEnsureToRepository(t *testing.T) {
accountID := int64(204)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{},
},
}}
err := (&adminServiceImpl{accountRepo: repo}).UpdateAccountExtra(context.Background(), accountID, map[string]any{
codexFingerprintModeExtraKey: "device",
codexFingerprintSeedExtraKey: userSuppliedCodexFingerprintSeed,
})
require.NoError(t, err)
require.Len(t, repo.updates[accountID], 1)
require.Equal(t, "device", repo.updates[accountID][0][codexFingerprintModeExtraKey])
require.NotContains(t, repo.updates[accountID][0], codexFingerprintSeedExtraKey)
}
func TestBulkUpdateAccountsDoesNotPrewriteCodexSeed(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{301, 302},
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: userSuppliedCodexFingerprintSeed,
},
})
require.NoError(t, err)
require.Equal(t, 2, result.Success)
require.Empty(t, repo.updates, "bulk enable must not loop through UpdateExtra before BulkUpdate")
require.Len(t, repo.bulkUpdates, 1)
require.True(t, repo.bulkUpdates[0].EnsureCodexFingerprintSeed)
require.Equal(t, "session", repo.bulkUpdates[0].Extra[codexFingerprintModeExtraKey])
require.NotContains(t, repo.bulkUpdates[0].Extra, codexFingerprintSeedExtraKey)
}
type codexSeedDuplicateRepo struct {
*upstreamBillingProbeAccountRepo
}
func (r *codexSeedDuplicateRepo) CreateWithAccountGroups(ctx context.Context, account *Account, _ []AccountGroup) error {
return r.Create(ctx, account)
}
func TestDuplicateAccountDoesNotCopyCodexFingerprintSeed(t *testing.T) {
ctx := context.Background()
repo := &codexSeedDuplicateRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: make(map[int64]*Account)}}
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
source := &Account{
Name: "source",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: testCodexFingerprintSeed,
},
}
require.NoError(t, repo.Create(ctx, source))
duplicate, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
require.NoError(t, err)
require.NotEqual(t, source.ID, duplicate.ID)
require.NotContains(t, duplicate.Extra, codexFingerprintSeedExtraKey)
require.Equal(t, "session", duplicate.Extra[codexFingerprintModeExtraKey])
}
func TestDuplicateCreatePathMintsFreshSeedWhenEligible(t *testing.T) {
extra, err := duplicateAccountExtra(map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: testCodexFingerprintSeed,
})
require.NoError(t, err)
account, err := buildAccountForCreate(&CreateAccountInput{
Name: "eligible-copy",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: extra,
}, extra)
require.NoError(t, err)
require.NotEqual(t, testCodexFingerprintSeed, requireValidCodexFingerprintSeed(t, account.Extra))
require.Equal(t, "session", account.Extra[codexFingerprintModeExtraKey])
}
func TestAccountServiceCreateAndUpdateCodexSeedLifecycle(t *testing.T) {
ctx := context.Background()
repo := &upstreamBillingProbeAccountRepo{accounts: make(map[int64]*Account)}
svc := NewAccountService(repo, nil)
created, err := svc.Create(ctx, CreateAccountRequest{
Name: "legacy-create",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
codexFingerprintModeExtraKey: "session",
codexFingerprintSeedExtraKey: userSuppliedCodexFingerprintSeed,
},
})
require.NoError(t, err)
createdSeed := requireValidCodexFingerprintSeed(t, created.Extra)
require.NotEqual(t, userSuppliedCodexFingerprintSeed, createdSeed)
updateSeed := userSuppliedCodexFingerprintSeed
updated, err := svc.Update(ctx, created.ID, UpdateAccountRequest{
Extra: &map[string]any{
codexFingerprintModeExtraKey: "full",
codexFingerprintSeedExtraKey: updateSeed,
},
})
require.NoError(t, err)
require.Equal(t, createdSeed, requireValidCodexFingerprintSeed(t, updated.Extra))
require.Equal(t, "full", updated.Extra[codexFingerprintModeExtraKey])
}
@@ -0,0 +1,96 @@
//go:build unit
package service
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
)
// shadowSkipTestRepo 是满足 AccountRepository 接口的最小 stub(只实现 GetByID)。
// 其他方法通过嵌入 nil 接口值满足编译,若被误调则 panic,便于发现意外调用路径。
type shadowSkipTestRepo struct {
AccountRepository
account *Account
}
func (r *shadowSkipTestRepo) GetByID(_ context.Context, id int64) (*Account, error) {
if r.account == nil || r.account.ID != id {
return nil, ErrAccountNotFound
}
return r.account, nil
}
func newShadowTestGinCtx() *gin.Context {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/200/test", nil)
return c
}
// --- 1. CanRefresh 守卫 ---
// TestOpenAITokenRefresherSkipsShadow 验证影子账号不被后台 token 刷新器处理。
func TestOpenAITokenRefresherSkipsShadow(t *testing.T) {
pid := int64(100)
r := NewOpenAITokenRefresher(nil, nil)
// 影子账号:ParentAccountID 非 nil → CanRefresh 应返回 false
require.False(t, r.CanRefresh(&Account{ID: 200, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &pid}))
// 普通账号:有 refresh_token → CanRefresh 应返回 true
require.True(t, r.CanRefresh(&Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"refresh_token": "RT"}}))
}
// --- 2. TestAccountConnection 影子凭据解析 ---
// TestAccountTestServiceSkipsShadow 验证影子账号连接测试不再早拒,而是尝试解析母账号凭据。
func TestAccountTestServiceSkipsShadow(t *testing.T) {
pid := int64(100)
shadow := &Account{
ID: 200,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
ParentAccountID: &pid,
}
repo := &shadowSkipTestRepo{account: shadow}
svc := &AccountTestService{accountRepo: repo}
c := newShadowTestGinCtx()
err := svc.TestAccountConnection(c, 200, "", "", "")
require.Error(t, err)
require.Contains(t, err.Error(), "resolve spark shadow parent")
}
// --- 3. EnsureOpenAIPrivacy 守卫 ---
// TestEnsureOpenAIPrivacySkipsShadow 验证影子账号跳过隐私设置(不调用 privacyClientFactory)。
// 影子账号透传母账号凭据,但 Extra 通常为空,需给它一个 access_token 才能让
// 现有的 token=="" 提前返回路径失效,从而真实验证 IsCredentialShadow 守卫。
func TestEnsureOpenAIPrivacySkipsShadow(t *testing.T) {
pid := int64(100)
shadow := &Account{
ID: 200,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
ParentAccountID: &pid,
// 提供 access_token:没有影子守卫时会进入 factory 调用
Credentials: map[string]any{"access_token": "shadow-passthrough-token"},
}
privacyCalled := false
svc := &adminServiceImpl{
privacyClientFactory: func(proxyURL string) (*req.Client, error) {
privacyCalled = true
return nil, errors.New("should not reach factory for shadow account")
},
}
got := svc.EnsureOpenAIPrivacy(context.Background(), shadow)
require.Equal(t, "", got)
require.False(t, privacyCalled, "privacyClientFactory 不应被影子账号触发")
}
@@ -0,0 +1,63 @@
package service
import (
"context"
"log/slog"
)
type accountCredentialsUpdater interface {
UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error
}
func persistAccountCredentials(ctx context.Context, repo AccountRepository, account *Account, credentials map[string]any) error {
if repo == nil || account == nil {
return nil
}
// 安全不变量:spark 影子账号恒不持凭据(凭据透传母账号)。这是凭据写入的唯一汇聚点
// (token 刷新 / 订阅补全 / CRS 创建后刷新等全部经此),在此对影子早返 no-op 是
// defense-in-depth——即便某条上游路径漏判,也不会把凭据落到影子行(外审第6轮 P1)。
if account.IsCredentialShadow() {
slog.Warn("skip persisting credentials to spark shadow account",
"account_id", account.ID, "parent_id", *account.ParentAccountID)
return nil
}
account.Credentials = shallowCopyMap(credentials)
if updater, ok := any(repo).(accountCredentialsUpdater); ok {
return updater.UpdateCredentials(ctx, account.ID, account.Credentials)
}
return repo.Update(ctx, account)
}
// sparkShadowAllowedCredentialKeys 是 spark 影子账号唯一可写的凭据键集合(仅模型映射)。
// 校验(isAllowed)与 sanitize 共用此单一来源,避免两处独立硬编码列表漂移。
var sparkShadowAllowedCredentialKeys = map[string]struct{}{
"model_mapping": {},
"compact_model_mapping": {},
}
func isAllowedSparkShadowCredentialsUpdate(credentials map[string]any) bool {
if credentials == nil {
return true
}
for key := range credentials {
if _, ok := sparkShadowAllowedCredentialKeys[key]; !ok {
return false
}
}
return true
}
func sanitizeSparkShadowCredentials(credentials map[string]any) map[string]any {
if len(credentials) == 0 {
return map[string]any{}
}
out := make(map[string]any, len(sparkShadowAllowedCredentialKeys))
for key := range sparkShadowAllowedCredentialKeys {
if value, ok := credentials[key]; ok && value != nil {
out[key] = value
}
}
return out
}
@@ -0,0 +1,52 @@
package service
// SensitiveCredentialKeys 列出 Account.Credentials JSON map 中绝不允许返回到前端的子键。
// dto 层做响应脱敏、service 层做更新合并都引用此清单——新增凭证类型时务必同步。
var SensitiveCredentialKeys = []string{
// OAuth
"access_token", "refresh_token", "id_token", "agent_private_key",
// API Key 类
"api_key", "session_key", "cookie",
// Grok Web SSO / password (must never persist or echo after Build OAuth)
"password", "sso_token", "sso", "sso-rw", "clearTextPassword",
// 云服务凭据
"aws_secret_access_key", "aws_session_token",
"service_account_json", "service_account", "private_key",
}
var sensitiveCredentialKeySet = func() map[string]struct{} {
m := make(map[string]struct{}, len(SensitiveCredentialKeys))
for _, k := range SensitiveCredentialKeys {
m[k] = struct{}{}
}
return m
}()
// IsSensitiveCredentialKey 判断指定键是否为敏感凭证子键。
func IsSensitiveCredentialKey(key string) bool {
_, ok := sensitiveCredentialKeySet[key]
return ok
}
// MergePreservingSensitiveCreds 把 incoming 写入 existing 之上,但敏感子键采用"incoming 没提供就保留 existing"
// 的语义。返回新的 map,不修改入参。
//
// 用途:前端编辑账号通常采用"全对象 PUT"模式;脱敏后前端 spread 旧 credentials 时不会带上敏感键,
// 直接覆盖会清空已有 token。此函数保证:
// - 非敏感键:完全由 incoming 决定(用户可以编辑、删除非敏感字段)。
// - 敏感键:incoming 显式提供则覆盖(用户主动旋转 token),否则保留 existing。
func MergePreservingSensitiveCreds(existing, incoming map[string]any) map[string]any {
out := make(map[string]any, len(incoming)+len(SensitiveCredentialKeys))
for k, v := range incoming {
out[k] = v
}
for _, key := range SensitiveCredentialKeys {
if _, hasIncoming := incoming[key]; hasIncoming {
continue
}
if existingVal, ok := existing[key]; ok {
out[key] = existingVal
}
}
return out
}
@@ -0,0 +1,90 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestMergePreservingSensitiveCreds_PreservesSensitiveWhenIncomingMissing(t *testing.T) {
existing := map[string]any{
"refresh_token": "rt-old",
"access_token": "at-old",
"api_key": "sk-old",
"base_url": "https://old.example.com",
}
incoming := map[string]any{
"base_url": "https://new.example.com",
"model_mapping": map[string]any{"foo": "bar"},
}
out := MergePreservingSensitiveCreds(existing, incoming)
require.Equal(t, "rt-old", out["refresh_token"], "incoming 没传 refresh_token,应保留 existing")
require.Equal(t, "at-old", out["access_token"])
require.Equal(t, "sk-old", out["api_key"])
require.Equal(t, "https://new.example.com", out["base_url"], "非敏感键由 incoming 决定")
require.Equal(t, map[string]any{"foo": "bar"}, out["model_mapping"])
}
func TestMergePreservingSensitiveCreds_OverwritesWhenIncomingProvidesSensitive(t *testing.T) {
existing := map[string]any{
"refresh_token": "rt-old",
"api_key": "sk-old",
}
incoming := map[string]any{
"refresh_token": "rt-new",
// 显式没传 api_key —— 应保留
}
out := MergePreservingSensitiveCreds(existing, incoming)
require.Equal(t, "rt-new", out["refresh_token"], "incoming 显式传入应覆盖")
require.Equal(t, "sk-old", out["api_key"], "incoming 没传应保留")
}
func TestMergePreservingSensitiveCreds_DoesNotMutateInputs(t *testing.T) {
existing := map[string]any{"refresh_token": "rt"}
incoming := map[string]any{"base_url": "x"}
_ = MergePreservingSensitiveCreds(existing, incoming)
require.Equal(t, "rt", existing["refresh_token"])
require.NotContains(t, existing, "base_url")
require.Equal(t, "x", incoming["base_url"])
require.NotContains(t, incoming, "refresh_token")
}
func TestMergePreservingSensitiveCreds_NilInputs(t *testing.T) {
out := MergePreservingSensitiveCreds(nil, map[string]any{"base_url": "x"})
require.Equal(t, "x", out["base_url"])
require.NotContains(t, out, "refresh_token")
out2 := MergePreservingSensitiveCreds(map[string]any{"refresh_token": "rt"}, nil)
require.Equal(t, "rt", out2["refresh_token"])
}
func TestMergePreservingSensitiveCreds_NonSensitiveDeletionAllowed(t *testing.T) {
existing := map[string]any{
"refresh_token": "rt",
"base_url": "https://old",
"project_id": "p1",
}
incoming := map[string]any{
"base_url": "https://new",
// 不带 project_id —— 等同删除(非敏感键由 incoming 决定)
}
out := MergePreservingSensitiveCreds(existing, incoming)
require.Equal(t, "rt", out["refresh_token"], "敏感键保留")
require.Equal(t, "https://new", out["base_url"])
require.NotContains(t, out, "project_id", "非敏感键 incoming 不传 = 删除")
}
func TestIsSensitiveCredentialKey(t *testing.T) {
require.True(t, IsSensitiveCredentialKey("refresh_token"))
require.True(t, IsSensitiveCredentialKey("api_key"))
require.True(t, IsSensitiveCredentialKey("private_key"))
require.False(t, IsSensitiveCredentialKey("base_url"))
require.False(t, IsSensitiveCredentialKey(""))
require.False(t, IsSensitiveCredentialKey("model_mapping"))
}
@@ -0,0 +1,71 @@
package service
import (
"context"
"log"
"sync"
"time"
)
// AccountExpiryService periodically pauses expired accounts when auto-pause is enabled.
type AccountExpiryService struct {
accountRepo AccountRepository
interval time.Duration
stopCh chan struct{}
stopOnce sync.Once
wg sync.WaitGroup
}
func NewAccountExpiryService(accountRepo AccountRepository, interval time.Duration) *AccountExpiryService {
return &AccountExpiryService{
accountRepo: accountRepo,
interval: interval,
stopCh: make(chan struct{}),
}
}
func (s *AccountExpiryService) Start() {
if s == nil || s.accountRepo == nil || s.interval <= 0 {
return
}
s.wg.Add(1)
go func() {
defer s.wg.Done()
ticker := time.NewTicker(s.interval)
defer ticker.Stop()
s.runOnce()
for {
select {
case <-ticker.C:
s.runOnce()
case <-s.stopCh:
return
}
}
}()
}
func (s *AccountExpiryService) Stop() {
if s == nil {
return
}
s.stopOnce.Do(func() {
close(s.stopCh)
})
s.wg.Wait()
}
func (s *AccountExpiryService) runOnce() {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
updated, err := s.accountRepo.AutoPauseExpiredAccounts(ctx, time.Now())
if err != nil {
log.Printf("[AccountExpiry] Auto pause expired accounts failed: %v", err)
return
}
if updated > 0 {
log.Printf("[AccountExpiry] Auto paused %d expired accounts", updated)
}
}
@@ -0,0 +1,65 @@
package service
import infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
// ValidateGrokMediaEligibilityExtra validates the optional per-account media
// routing override. A nil value removes the override and restores automatic
// provider-observation routing.
func ValidateGrokMediaEligibilityExtra(platform string, extra map[string]any) error {
if platform != PlatformGrok || extra == nil {
return nil
}
raw, exists := extra[GrokMediaEligibleExtraKey]
if !exists || raw == nil {
return nil
}
if _, ok := raw.(bool); !ok {
return infraerrors.BadRequest("GROK_MEDIA_ELIGIBILITY_INVALID", "grok_media_eligible must be a boolean or null")
}
return nil
}
func normalizeGrokMediaEligibilityExtra(platform string, extra map[string]any) (map[string]any, error) {
if platform != PlatformGrok {
return extra, nil
}
if err := ValidateGrokMediaEligibilityExtra(platform, extra); err != nil {
return nil, err
}
if extra == nil {
return nil, nil
}
normalized := shallowCopyMap(extra)
if normalized[GrokMediaEligibleExtraKey] == nil {
delete(normalized, GrokMediaEligibleExtraKey)
}
return normalized, nil
}
func normalizeGrokMediaEligibilityUpdateExtra(account *Account, input *UpdateAccountInput, normalized map[string]any) (map[string]any, error) {
if account == nil || account.Platform != PlatformGrok {
return normalized, nil
}
if input == nil {
return nil, infraerrors.BadRequest("INVALID_ACCOUNT_INPUT", "account update input is required")
}
if err := ValidateGrokMediaEligibilityExtra(account.Platform, input.Extra); err != nil {
return nil, err
}
if normalized == nil {
normalized = make(map[string]any)
} else {
normalized = shallowCopyMap(normalized)
}
raw, provided := input.Extra[GrokMediaEligibleExtraKey]
if provided {
if raw == nil {
delete(normalized, GrokMediaEligibleExtraKey)
}
return normalized, nil
}
if current, ok := account.Extra[GrokMediaEligibleExtraKey].(bool); ok {
normalized[GrokMediaEligibleExtraKey] = current
}
return normalized, nil
}
@@ -0,0 +1,180 @@
//go:build unit
package service
import (
"context"
"net/http"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
func TestGrokMediaGenerationEligibility(t *testing.T) {
weeklyUsagePercent := 12.5
forbiddenBilling := &xai.BillingSummary{
StatusCode: http.StatusForbidden,
WeeklyStatusCode: http.StatusForbidden,
MonthlyStatusCode: http.StatusForbidden,
}
weeklyAllowance := &xai.BillingSummary{
PeriodType: "weekly",
UsagePercent: &weeklyUsagePercent,
StatusCode: http.StatusOK,
WeeklyStatusCode: http.StatusOK,
}
weeklyForbidden := &xai.BillingSummary{
StatusCode: http.StatusOK,
WeeklyStatusCode: http.StatusForbidden,
MonthlyStatusCode: http.StatusOK,
}
monthlyForbidden := &xai.BillingSummary{
StatusCode: http.StatusOK,
WeeklyStatusCode: http.StatusOK,
MonthlyStatusCode: http.StatusForbidden,
}
tests := []struct {
name string
account *Account
want bool
wantReason string
}{
{name: "nil account", account: nil, want: false, wantReason: "not_grok"},
{name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"},
{name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"},
{name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"},
{name: "weekly paid usage is eligible without inferring from period type", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"},
{name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"},
{name: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"},
{name: "monthly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: monthlyForbidden}}, want: false, wantReason: "billing_forbidden"},
{name: "malformed billing observation fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: false, wantReason: "billing_unobserved"},
{name: "malformed override falls back to observations", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "false", grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"},
{name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"},
{name: "explicit enable wins over forbidden probe", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: forbiddenBilling}}, want: true, wantReason: "override_enabled"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, reason := tt.account.GrokMediaGenerationEligibility()
require.Equal(t, tt.want, got)
require.Equal(t, tt.wantReason, reason)
})
}
}
func TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.T) {
unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}
eligible, reason := unobserved.GrokMediaGenerationEligibility()
require.False(t, eligible)
require.Equal(t, "billing_unobserved", reason)
require.True(t, unobserved.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
inconclusive := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{
StatusCode: http.StatusOK,
Partial: true,
}},
}
require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
}
func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) {
account := &Account{
ID: 1,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Extra: map[string]any{GrokMediaEligibleExtraKey: false},
}
require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
require.False(t, isOpenAICompatibleAccountEligibleForRequest(
context.Background(), account, PlatformGrok, "grok-imagine-video", false,
OpenAIEndpointCapabilityGrokMediaGeneration,
))
}
func TestNormalizeGrokMediaEligibilityExtra(t *testing.T) {
t.Run("boolean override is accepted", func(t *testing.T) {
extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: false})
require.NoError(t, err)
require.Equal(t, false, extra[GrokMediaEligibleExtraKey])
})
t.Run("null clears override", func(t *testing.T) {
extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: nil})
require.NoError(t, err)
require.NotContains(t, extra, GrokMediaEligibleExtraKey)
})
t.Run("malformed override is rejected", func(t *testing.T) {
_, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: "false"})
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
})
t.Run("other platforms ignore provider owned value", func(t *testing.T) {
extra := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
normalized, err := normalizeGrokMediaEligibilityExtra(PlatformOpenAI, extra)
require.NoError(t, err)
require.Equal(t, extra, normalized)
})
}
func TestNormalizeGrokMediaEligibilityUpdateExtra(t *testing.T) {
account := &Account{Platform: PlatformGrok, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}
t.Run("omitted override preserves current value", func(t *testing.T) {
input := &UpdateAccountInput{Extra: map[string]any{"quota_used": float64(1)}}
normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{"quota_used": float64(1)})
require.NoError(t, err)
require.Equal(t, false, normalized[GrokMediaEligibleExtraKey])
})
t.Run("null removes current override", func(t *testing.T) {
input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: nil}}
normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: nil})
require.NoError(t, err)
require.NotContains(t, normalized, GrokMediaEligibleExtraKey)
require.Contains(t, input.Extra, GrokMediaEligibleExtraKey)
})
t.Run("provided boolean replaces current override", func(t *testing.T) {
input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: true}}
normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: true})
require.NoError(t, err)
require.Equal(t, true, normalized[GrokMediaEligibleExtraKey])
})
t.Run("malformed override is rejected on update", func(t *testing.T) {
input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "false"}}
_, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, nil)
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
})
t.Run("non grok update is unchanged", func(t *testing.T) {
input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}}
normalized := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
got, err := normalizeGrokMediaEligibilityUpdateExtra(&Account{Platform: PlatformOpenAI}, input, normalized)
require.NoError(t, err)
require.Equal(t, normalized, got)
})
}
+13
View File
@@ -0,0 +1,13 @@
package service
import "time"
type AccountGroup struct {
AccountID int64
GroupID int64
Priority int
CreatedAt time.Time
Account *Account
Group *Group
}
@@ -0,0 +1,292 @@
package service
import (
"net/http"
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"golang.org/x/net/http/httpguts"
)
// 请求头覆写(header override):对 Anthropic / OpenAI / Kimi / Zhipu / DeepSeek
// 平台的 api_key 账号,以及 Grok 平台的 api_key / oauth 账号生效。
// 管理员在账号上配置一组 header name -> value,转发到上游前用配置值覆盖同名请求头
// (匹配不区分大小写);value 为空的条目视为"未填写",不参与覆盖。
const (
credKeyHeaderOverrideEnabled = "header_override_enabled"
credKeyHeaderOverrides = "header_overrides"
maxHeaderOverrideEntries = 64
maxHeaderOverrideNameLength = 200
maxHeaderOverrideValueLength = 8192
)
// headerOverrideBlockedNames 禁止覆写的请求头(小写)。
// - 连接控制/逐跳头:由 HTTP 栈管理,覆写会破坏请求传输;
// - host/content-length:由 Go 的 Request.Host / ContentLength 字段管理,header 覆写不生效或产生冲突;
// - content-type:承载报文框架信息(multipart boundary 为每请求随机值),静态覆写必然与 body 不匹配;
// - authorization/x-api-key/cookie 等:上游认证头由账号凭据统一注入,禁止通过覆写篡改或重新引入;
// - accept-encoding:强制压缩会破坏网关对上游流式响应(SSE/usage)的解析;
// - sec-websocket-*WebSocket 握手头由拨号器管理(OpenAI WS 模式);
// - session_id/x-claude-code-session-id/x-grok-conv-id 等:逐请求会话隔离头,
// 固定值会造成会话串扰。
var headerOverrideBlockedNames = map[string]struct{}{
"host": {},
"content-length": {},
"content-type": {},
"transfer-encoding": {},
"connection": {},
"keep-alive": {},
"proxy-authenticate": {},
"proxy-authorization": {},
"proxy-connection": {},
"te": {},
"trailer": {},
"upgrade": {},
"authorization": {},
"x-api-key": {},
"x-goog-api-key": {},
"cookie": {},
"accept-encoding": {},
"sec-websocket-key": {},
"sec-websocket-version": {},
"sec-websocket-extensions": {},
"sec-websocket-protocol": {},
"sec-websocket-accept": {},
"session_id": {},
"conversation_id": {},
"x-codex-turn-state": {},
"x-codex-turn-metadata": {},
"chatgpt-account-id": {},
"x-claude-code-session-id": {},
"x-client-request-id": {},
"x-grok-conv-id": {},
}
func isHeaderOverrideBlockedName(lowerName string) bool {
_, blocked := headerOverrideBlockedNames[lowerName]
return blocked
}
// IsHeaderOverrideEligible 报告账号类型是否支持请求头覆写。
// Anthropic / OpenAI / Kimi / Zhipu / DeepSeek 仅开放 api_key 账号;
// Grok 额外开放 oauth 账号——
// 订阅流量改发自定义转发地址时,通常需要补充中间层要求的准入头。
func (a *Account) IsHeaderOverrideEligible() bool {
if a == nil {
return false
}
switch a.Platform {
case PlatformAnthropic, PlatformOpenAI, PlatformKimi, PlatformZhipu, PlatformDeepseek:
return a.Type == AccountTypeAPIKey
case PlatformGrok:
return a.Type == AccountTypeAPIKey || a.Type == AccountTypeOAuth
default:
return false
}
}
// IsHeaderOverrideEnabled 报告账号是否启用了请求头覆写。
func (a *Account) IsHeaderOverrideEnabled() bool {
if !a.IsHeaderOverrideEligible() || a.Credentials == nil {
return false
}
enabled, ok := a.Credentials[credKeyHeaderOverrideEnabled].(bool)
return ok && enabled
}
// GetHeaderOverrides 返回生效的请求头覆写表(key 统一小写)。
// 未启用、不符合平台/类型条件或配置为空时返回 nil。
// 空 value 的条目(模板占位)与非法/禁止的 header 名会被跳过。
// 结果带热路径缓存(同 GetModelMapping 先例):同一 credentials 映射在
// 一次请求 / 一条 WS 会话内的多次调用只做一次解析与校验。
func (a *Account) GetHeaderOverrides() map[string]string {
if !a.IsHeaderOverrideEnabled() {
return nil
}
rawMapping, rawIsAnyMap := a.Credentials[credKeyHeaderOverrides].(map[string]any)
if !rawIsAnyMap {
// 非 JSON 反序列化产物(如直接注入的 map[string]string):直接解析,不缓存
return resolveHeaderOverrides(stringMappingFromRaw(a.Credentials[credKeyHeaderOverrides]))
}
credentialsPtr := mapPtr(a.Credentials)
rawPtr := mapPtr(rawMapping)
rawLen := len(rawMapping)
rawSig := uint64(0)
rawSigReady := false
if a.headerOverrideCacheReady &&
a.headerOverrideCacheCredentialsPtr == credentialsPtr &&
a.headerOverrideCacheRawPtr == rawPtr &&
a.headerOverrideCacheRawLen == rawLen {
rawSig = modelMappingSignature(rawMapping)
rawSigReady = true
if a.headerOverrideCacheRawSig == rawSig {
return a.headerOverrideCache
}
}
overrides := resolveHeaderOverrides(stringMappingFromRaw(rawMapping))
if !rawSigReady {
rawSig = modelMappingSignature(rawMapping)
}
a.headerOverrideCache = overrides
a.headerOverrideCacheReady = true
a.headerOverrideCacheCredentialsPtr = credentialsPtr
a.headerOverrideCacheRawPtr = rawPtr
a.headerOverrideCacheRawLen = rawLen
a.headerOverrideCacheRawSig = rawSig
return overrides
}
// resolveHeaderOverrides 解析并防御性过滤原始覆写表:保存路径已做校验,
// 这里兜底未经 Normalize 落库的数据(含名单扩充前保存的旧配置),非法条目直接跳过。
func resolveHeaderOverrides(raw map[string]string) map[string]string {
if len(raw) == 0 {
return nil
}
result := make(map[string]string, len(raw))
for name, value := range raw {
lowerName, value, err := normalizeHeaderOverrideEntry(name, value)
if err != nil || lowerName == "" || value == "" {
continue
}
result[lowerName] = value
}
if len(result) == 0 {
return nil
}
return result
}
// HeaderOverrideValue 返回指定 header(小写名)的生效覆写值。
// 供转发链路在 header 写入前感知覆写结果(如 anthropic-beta 需要参与 body 净化)。
func (a *Account) HeaderOverrideValue(lowerName string) (string, bool) {
value, ok := a.GetHeaderOverrides()[lowerName]
return value, ok
}
// ApplyHeaderOverrides 将账号配置的请求头覆写应用到出站请求头。
// 对每个覆写条目:先删除所有大小写变体(转发链路会以 wire casing 直接写入 map
// 可能存在非 canonical key),再按已知 wire casing 写入,避免产生重复头。
// 账号未启用或不符合条件时为 no-op,可安全地在 OAuth/api_key 共用的构建器中调用。
func (a *Account) ApplyHeaderOverrides(h http.Header) {
if h == nil {
return
}
overrides := a.GetHeaderOverrides()
if len(overrides) == 0 {
return
}
// 覆写名两两不同(大小写不敏感)且各自只操作同名键,应用顺序不影响结果。
// 全量 EqualFold 扫描兜底删除任意 casing 的既有键:透传链路可能保留客户端
// 原始 casing,非 canonical/wire casing 的键 deleteHeaderAllForms 覆盖不到。
for name, value := range overrides {
for existing := range h {
if strings.EqualFold(existing, name) {
delete(h, existing)
}
}
h[resolveWireCasing(name)] = []string{value}
}
}
// NormalizeHeaderOverrideCredentials 校验并原地规范化 credentials 中的请求头覆写字段。
// 供账号创建/更新/批量更新的保存路径调用;credentials 未携带相关字段时为 no-op。
// 规范化内容:header 名转小写并去除首尾空白,value 去除首尾空白,丢弃名和值均为空的条目。
func NormalizeHeaderOverrideCredentials(credentials map[string]any) error {
if credentials == nil {
return nil
}
if raw, ok := credentials[credKeyHeaderOverrideEnabled]; ok && raw != nil {
if _, isBool := raw.(bool); !isBool {
return infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header_override_enabled must be a boolean")
}
}
raw, ok := credentials[credKeyHeaderOverrides]
if !ok || raw == nil {
return nil
}
var entries map[string]any
switch m := raw.(type) {
case map[string]any:
entries = m
case map[string]string:
entries = make(map[string]any, len(m))
for k, v := range m {
entries[k] = v
}
default:
return infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header_overrides must be an object of header name to string value")
}
if len(entries) > maxHeaderOverrideEntries {
return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header_overrides supports at most %d entries", maxHeaderOverrideEntries)
}
normalized := make(map[string]any, len(entries))
for name, rawValue := range entries {
value, isString := rawValue.(string)
if !isString {
return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header %q value must be a string", name)
}
lowerName, value, err := normalizeHeaderOverrideEntry(name, value)
if err != nil {
return err
}
if lowerName == "" {
continue // 丢弃完全为空的占位行
}
if _, dup := normalized[lowerName]; dup {
return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"duplicate header name %q (matching is case-insensitive)", lowerName)
}
normalized[lowerName] = value
}
credentials[credKeyHeaderOverrides] = normalized
return nil
}
// normalizeHeaderOverrideEntry 校验并规范化单个覆写条目,保存路径(Normalizeerr → 400
// 与应用路径(resolveHeaderOverrideserr → 跳过)共用同一套规则,避免两处校验漂移。
// 名和值均为空表示空占位行,返回 ("", "", nil);空 value 的具名条目合法(模板占位)。
func normalizeHeaderOverrideEntry(name, value string) (string, string, error) {
lowerName := strings.ToLower(strings.TrimSpace(name))
value = strings.TrimSpace(value)
if lowerName == "" {
if value == "" {
return "", "", nil
}
return "", "", infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header name must not be empty")
}
if len(lowerName) > maxHeaderOverrideNameLength {
return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header name %q exceeds %d characters", lowerName, maxHeaderOverrideNameLength)
}
if !httpguts.ValidHeaderFieldName(lowerName) {
return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"invalid header name %q", lowerName)
}
if isHeaderOverrideBlockedName(lowerName) {
return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header %q is not allowed to be overridden", lowerName)
}
if len(value) > maxHeaderOverrideValueLength {
return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header %q value exceeds %d characters", lowerName, maxHeaderOverrideValueLength)
}
if !httpguts.ValidHeaderFieldValue(value) {
return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE",
"header %q has an invalid value", lowerName)
}
return lowerName, value, nil
}
@@ -0,0 +1,346 @@
//go:build unit
package service
import (
"net/http"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func headerOverrideTestAccount(platform, accountType string, credentials map[string]any) *Account {
return &Account{
Platform: platform,
Type: accountType,
Credentials: credentials,
}
}
func TestIsHeaderOverrideEligible(t *testing.T) {
tests := []struct {
name string
platform string
accType string
want bool
}{
{"anthropic apikey", PlatformAnthropic, AccountTypeAPIKey, true},
{"openai apikey", PlatformOpenAI, AccountTypeAPIKey, true},
{"kimi apikey", PlatformKimi, AccountTypeAPIKey, true},
{"zhipu apikey", PlatformZhipu, AccountTypeAPIKey, true},
{"deepseek apikey", PlatformDeepseek, AccountTypeAPIKey, true},
{"anthropic oauth", PlatformAnthropic, AccountTypeOAuth, false},
{"openai oauth", PlatformOpenAI, AccountTypeOAuth, false},
{"kimi oauth", PlatformKimi, AccountTypeOAuth, false},
{"zhipu oauth", PlatformZhipu, AccountTypeOAuth, false},
{"deepseek oauth", PlatformDeepseek, AccountTypeOAuth, false},
{"gemini apikey", PlatformGemini, AccountTypeAPIKey, false},
{"grok apikey", PlatformGrok, AccountTypeAPIKey, true},
{"grok oauth", PlatformGrok, AccountTypeOAuth, true},
{"antigravity apikey", PlatformAntigravity, AccountTypeAPIKey, false},
{"anthropic bedrock", PlatformAnthropic, AccountTypeBedrock, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
acc := headerOverrideTestAccount(tt.platform, tt.accType, nil)
require.Equal(t, tt.want, acc.IsHeaderOverrideEligible())
})
}
var nilAccount *Account
require.False(t, nilAccount.IsHeaderOverrideEligible())
require.False(t, nilAccount.IsHeaderOverrideEnabled())
require.Nil(t, nilAccount.GetHeaderOverrides())
}
func TestIsHeaderOverrideEnabled(t *testing.T) {
acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
})
require.True(t, acc.IsHeaderOverrideEnabled())
// 未配置 / 非 bool / false 均视为未启用
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, nil).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: "true",
}).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: false,
}).IsHeaderOverrideEnabled())
// 不符合平台/类型条件时即使配置了 true 也不启用
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{
credKeyHeaderOverrideEnabled: true,
}).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformGemini, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
}).IsHeaderOverrideEnabled())
}
func TestGetHeaderOverrides(t *testing.T) {
acc := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"User-Agent": "my-agent/1.0", // 大写 key 归一化为小写
" X-App ": "cli", // 名称去空白
"x-empty": "", // 空 value(模板占位)跳过
"authorization": "Bearer leaked", // 禁止覆写的头跳过
"bad name": "value", // 非法 header 名跳过
"x-padded": " padded ", // value 去空白
},
})
overrides := acc.GetHeaderOverrides()
require.Equal(t, map[string]string{
"user-agent": "my-agent/1.0",
"x-app": "cli",
"x-padded": "padded",
}, overrides)
// 未启用时返回 nil
disabled := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrides: map[string]any{"user-agent": "x"},
})
require.Nil(t, disabled.GetHeaderOverrides())
// 启用但全部为空 value 时返回 nil
empty := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"user-agent": ""},
})
require.Nil(t, empty.GetHeaderOverrides())
// 未经 Normalize 落库的超长数据 / WebSocket 握手头在应用时被防御性跳过
oversizedValue := strings.Repeat("a", maxHeaderOverrideValueLength+1)
defensive := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"x-big": oversizedValue,
"sec-websocket-key": "forged",
"content-type": "application/json", // 名单扩充前落库的数据也要被拦截
"x-claude-code-session-id": "pinned-session",
"x-ok": "ok",
},
})
require.Equal(t, map[string]string{"x-ok": "ok"}, defensive.GetHeaderOverrides())
}
func TestApplyHeaderOverrides(t *testing.T) {
acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"user-agent": "override-agent/2.0",
"anthropic-beta": "custom-beta-1",
"x-custom": "custom-value",
},
})
h := http.Header{}
// 模拟转发链路:canonical key 与 wire casing 原样 key 混合存在
h.Set("User-Agent", "claude-cli/2.1.161 (external, cli)")
h["anthropic-beta"] = []string{"claude-code-20250219,oauth-2025-04-20"} // 非 canonical 原样 key
h.Set("Content-Type", "application/json")
acc.ApplyHeaderOverrides(h)
// user-agent 覆盖且只有一个值(已知头恢复 wire casing
require.Equal(t, []string{"override-agent/2.0"}, h["User-Agent"])
// anthropic-beta:非 canonical 旧值被清除,写入 wire casing(小写)
require.Equal(t, []string{"custom-beta-1"}, h["anthropic-beta"])
require.Empty(t, h["Anthropic-Beta"])
// 新增头(未知头以小写原样键写入,与转发链路 wire casing 约定一致)
require.Equal(t, []string{"custom-value"}, h["x-custom"])
require.Equal(t, "custom-value", getHeaderRaw(h, "x-custom"))
// 未覆写的头不受影响
require.Equal(t, "application/json", h.Get("Content-Type"))
// 覆盖后不存在任何大小写重复
count := 0
for k := range h {
if k == "anthropic-beta" || k == "Anthropic-Beta" {
count++
}
}
require.Equal(t, 1, count)
}
func TestApplyHeaderOverridesNoOpPaths(t *testing.T) {
baseline := func() http.Header {
h := http.Header{}
h.Set("User-Agent", "orig")
return h
}
// OAuth 账号:即使配置了覆写也不生效
oauth := headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"},
})
h := baseline()
oauth.ApplyHeaderOverrides(h)
require.Equal(t, "orig", h.Get("User-Agent"))
// 未启用开关
off := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"},
})
h = baseline()
off.ApplyHeaderOverrides(h)
require.Equal(t, "orig", h.Get("User-Agent"))
// 禁止覆写的头(authorization / x-api-key / host 等)不会被应用
blocked := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"Authorization": "Bearer evil",
"X-Api-Key": "evil",
"Host": "evil.example.com",
"Content-Length": "0",
},
})
h = http.Header{}
h.Set("Authorization", "Bearer real-key")
blocked.ApplyHeaderOverrides(h)
require.Equal(t, "Bearer real-key", h.Get("Authorization"))
require.Empty(t, h.Get("X-Api-Key"))
require.Empty(t, h.Get("Host"))
// nil header 不 panic
blocked.ApplyHeaderOverrides(nil)
}
func TestNormalizeHeaderOverrideCredentials(t *testing.T) {
t.Run("nil credentials no-op", func(t *testing.T) {
require.NoError(t, NormalizeHeaderOverrideCredentials(nil))
})
t.Run("missing keys no-op", func(t *testing.T) {
creds := map[string]any{"api_key": "sk-xxx"}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
_, exists := creds[credKeyHeaderOverrides]
require.False(t, exists)
})
t.Run("normalizes names and values", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
" User-Agent ": " my-agent ",
"X-App": "",
"": "", // 完全空行被丢弃
},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{
"user-agent": "my-agent",
"x-app": "",
}, creds[credKeyHeaderOverrides])
})
t.Run("accepts map[string]string input", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrides: map[string]string{"X-App": "cli"},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{"x-app": "cli"}, creds[credKeyHeaderOverrides])
})
t.Run("rejects non-bool enabled", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrideEnabled: "yes",
})
require.Error(t, err)
})
t.Run("rejects non-object overrides", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: []any{"user-agent"},
})
require.Error(t, err)
})
t.Run("rejects non-string value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": 123},
})
require.Error(t, err)
})
t.Run("rejects invalid header name", func(t *testing.T) {
for _, name := range []string{"bad name", "bad:name", "bad\nname", "值"} {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{name: "v"},
})
require.Error(t, err, "name %q should be rejected", name)
}
})
t.Run("rejects empty name with value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{" ": "v"},
})
require.Error(t, err)
})
t.Run("rejects blocked headers", func(t *testing.T) {
for _, name := range []string{
"Authorization", "x-api-key", "Host", "content-length", "Transfer-Encoding",
"connection", "accept-encoding", "Sec-WebSocket-Key", "session_id",
"conversation_id", "x-codex-turn-state", "chatgpt-account-id",
"Content-Type", "Cookie", "x-goog-api-key",
"X-Claude-Code-Session-Id", "x-client-request-id",
} {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{name: "v"},
})
require.Error(t, err, "blocked header %q should be rejected", name)
}
})
t.Run("allows tab inside value", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": "a\tb"},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{"x-app": "a\tb"}, creds[credKeyHeaderOverrides])
})
t.Run("rejects invalid value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": "bad\nvalue"},
})
require.Error(t, err)
})
t.Run("rejects duplicate names case-insensitively", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{
"User-Agent": "a",
"user-agent": "b",
},
})
require.Error(t, err)
})
t.Run("rejects too many entries", func(t *testing.T) {
entries := make(map[string]any, maxHeaderOverrideEntries+1)
for i := 0; i <= maxHeaderOverrideEntries; i++ {
entries["x-h-"+string(rune('a'+i%26))+string(rune('a'+(i/26)%26))+string(rune('a'+(i/676)%26))] = "v"
}
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: entries,
})
require.Error(t, err)
})
t.Run("rejects oversized value", func(t *testing.T) {
big := make([]byte, maxHeaderOverrideValueLength+1)
for i := range big {
big[i] = 'a'
}
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": string(big)},
})
require.Error(t, err)
})
}
@@ -0,0 +1,66 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAccount_IsInterceptWarmupEnabled(t *testing.T) {
tests := []struct {
name string
credentials map[string]any
expected bool
}{
{
name: "nil credentials",
credentials: nil,
expected: false,
},
{
name: "empty map",
credentials: map[string]any{},
expected: false,
},
{
name: "field not present",
credentials: map[string]any{"access_token": "tok"},
expected: false,
},
{
name: "field is true",
credentials: map[string]any{"intercept_warmup_requests": true},
expected: true,
},
{
name: "field is false",
credentials: map[string]any{"intercept_warmup_requests": false},
expected: false,
},
{
name: "field is string true",
credentials: map[string]any{"intercept_warmup_requests": "true"},
expected: false,
},
{
name: "field is int 1",
credentials: map[string]any{"intercept_warmup_requests": 1},
expected: false,
},
{
name: "field is nil",
credentials: map[string]any{"intercept_warmup_requests": nil},
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := &Account{Credentials: tt.credentials}
result := a.IsInterceptWarmupEnabled()
require.Equal(t, tt.expected, result)
})
}
}
@@ -0,0 +1,46 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func intPtrHelper(v int) *int { return &v }
func TestEffectiveLoadFactor_NilAccount(t *testing.T) {
var a *Account
require.Equal(t, 1, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_NilLoadFactor_PositiveConcurrency(t *testing.T) {
a := &Account{Concurrency: 5}
require.Equal(t, 5, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_NilLoadFactor_ZeroConcurrency(t *testing.T) {
a := &Account{Concurrency: 0}
require.Equal(t, 1, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_PositiveLoadFactor(t *testing.T) {
a := &Account{Concurrency: 5, LoadFactor: intPtrHelper(20)}
require.Equal(t, 20, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_ZeroLoadFactor_FallbackToConcurrency(t *testing.T) {
a := &Account{Concurrency: 5, LoadFactor: intPtrHelper(0)}
require.Equal(t, 5, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_NegativeLoadFactor_FallbackToConcurrency(t *testing.T) {
a := &Account{Concurrency: 3, LoadFactor: intPtrHelper(-1)}
require.Equal(t, 3, a.EffectiveLoadFactor())
}
func TestEffectiveLoadFactor_ZeroLoadFactor_ZeroConcurrency(t *testing.T) {
a := &Account{Concurrency: 0, LoadFactor: intPtrHelper(0)}
require.Equal(t, 1, a.EffectiveLoadFactor())
}
@@ -0,0 +1,292 @@
//go:build unit
package service
import (
"context"
"net/http"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
)
func TestAccountIsOpenAILongContextBillingEnabled(t *testing.T) {
tests := []struct {
name string
account *Account
want bool
}{
{name: "nil account is disabled", account: nil, want: false},
{name: "non OpenAI account is disabled", account: &Account{Platform: PlatformGrok}, want: false},
{name: "missing extra defaults disabled", account: &Account{Platform: PlatformOpenAI}, want: false},
{name: "missing key defaults disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{}}, want: false},
{name: "explicit true is enabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": true}}, want: true},
{name: "explicit false is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": false}}, want: false},
{name: "malformed value is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": "false"}}, want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, tt.account.IsOpenAILongContextBillingEnabled())
})
}
}
func TestNormalizeOpenAILongContextBillingExtra(t *testing.T) {
t.Run("OpenAI missing key persists disabled default", func(t *testing.T) {
extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, nil)
require.NoError(t, err)
require.Equal(t, false, extra["openai_long_context_billing_enabled"])
})
t.Run("OpenAI explicit false is preserved", func(t *testing.T) {
extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": false})
require.NoError(t, err)
require.Equal(t, false, extra["openai_long_context_billing_enabled"])
})
t.Run("OpenAI malformed value is rejected", func(t *testing.T) {
_, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": "false"})
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
})
t.Run("non OpenAI extra is unchanged", func(t *testing.T) {
extra, err := normalizeOpenAILongContextBillingExtra(PlatformGrok, nil)
require.NoError(t, err)
require.Nil(t, extra)
})
t.Run("non OpenAI malformed value is ignored", func(t *testing.T) {
extra := map[string]any{openAILongContextBillingEnabledKey: "provider-owned"}
normalized, err := normalizeOpenAILongContextBillingExtra(PlatformAnthropic, extra)
require.NoError(t, err)
require.Equal(t, extra, normalized)
})
}
type longContextBillingRepoStub struct {
accountRepoStub
account *Account
accounts []*Account
createdAccount *Account
updateExtraCalls int
bulkUpdateCalls int
}
func (r *longContextBillingRepoStub) Create(_ context.Context, account *Account) error {
account.ID = 1
r.account = account
r.createdAccount = account
return nil
}
func (r *longContextBillingRepoStub) GetByID(_ context.Context, _ int64) (*Account, error) {
return r.account, nil
}
func (r *longContextBillingRepoStub) GetByIDs(_ context.Context, _ []int64) ([]*Account, error) {
if r.accounts != nil {
return r.accounts, nil
}
if r.account == nil {
return nil, nil
}
return []*Account{r.account}, nil
}
func (r *longContextBillingRepoStub) Update(_ context.Context, account *Account) error {
r.account = account
return nil
}
func (r *longContextBillingRepoStub) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error {
r.updateExtraCalls++
return nil
}
func (r *longContextBillingRepoStub) BulkUpdate(_ context.Context, _ []int64, _ AccountBulkUpdate) (int64, error) {
r.bulkUpdateCalls++
return 1, nil
}
func TestAdminServiceCreateAccountDefaultsOpenAILongContextBillingDisabled(t *testing.T) {
repo := &longContextBillingRepoStub{}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
Name: "openai-account",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "test"},
SkipDefaultGroupBind: true,
})
require.NoError(t, err)
require.Same(t, account, repo.createdAccount)
require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey])
}
func TestAdminServiceCreateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
repo := &longContextBillingRepoStub{}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
Platform: PlatformOpenAI,
Extra: map[string]any{openAILongContextBillingEnabledKey: "false"},
})
require.Nil(t, account)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Nil(t, repo.createdAccount)
}
func TestAdminServiceUpdateAccountPreservesOpenAILongContextBillingOptOutWhenOmitted(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{openAILongContextBillingEnabledKey: false},
}}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{}})
require.NoError(t, err)
require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey])
}
func TestAdminServiceUpdateAccountAllowsExplicitCodexImportOptIn(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{"access_token": "old-token"},
Extra: map[string]any{
openAILongContextBillingEnabledKey: false,
"import_source": "codex_session",
},
}}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{
Credentials: map[string]any{"access_token": "new-token"},
Extra: map[string]any{
openAILongContextBillingEnabledKey: true,
"import_source": "codex_session",
},
})
require.NoError(t, err)
require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey])
}
func TestAdminServiceUpdateAccountAllowsExplicitOptInOutsideCodexImport(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
openAILongContextBillingEnabledKey: false,
"import_source": "codex_session",
},
}}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{
openAILongContextBillingEnabledKey: true,
"import_source": "codex_session",
}})
require.NoError(t, err)
require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey])
}
func TestAdminServiceUpdateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
svc := &adminServiceImpl{accountRepo: repo}
account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{
openAILongContextBillingEnabledKey: 1,
}})
require.Nil(t, account)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
}
func TestAdminServiceUpdateAccountExtraRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
svc := &adminServiceImpl{accountRepo: repo}
err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{
openAILongContextBillingEnabledKey: "true",
})
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Zero(t, repo.updateExtraCalls)
}
func TestAdminServiceUpdateAccountExtraAllowsProviderOwnedValueForNonOpenAIAccount(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformAnthropic}}
svc := &adminServiceImpl{accountRepo: repo}
err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{
openAILongContextBillingEnabledKey: "provider-owned",
})
require.NoError(t, err)
require.Equal(t, 1, repo.updateExtraCalls)
}
func TestAdminServiceBulkUpdateAccountsRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{openAILongContextBillingEnabledKey: []bool{true}},
})
require.Nil(t, result)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Zero(t, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccountsRejectsOpenAILongContextKeyForNonOpenAIAccounts(t *testing.T) {
repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformGrok}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
})
require.Nil(t, result)
var appErr *infraerrors.ApplicationError
require.ErrorAs(t, err, &appErr)
require.Equal(t, "OPENAI_BULK_TARGET_INVALID", appErr.Reason)
require.Zero(t, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccountsRejectsMalformedValueForMixedTargetsIncludingOpenAI(t *testing.T) {
repo := &longContextBillingRepoStub{accounts: []*Account{
{ID: 1, Platform: PlatformGrok},
{ID: 2, Platform: PlatformOpenAI},
}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
Extra: map[string]any{openAILongContextBillingEnabledKey: "malformed"},
})
require.Nil(t, result)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Zero(t, repo.bulkUpdateCalls)
}
@@ -0,0 +1,369 @@
package service
import "testing"
func TestAccountGetOpenAICompactMode(t *testing.T) {
tests := []struct {
name string
account *Account
want string
}{
{
name: "nil account defaults to auto",
want: OpenAICompactModeAuto,
},
{
name: "non openai account defaults to auto",
account: &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOn},
},
want: OpenAICompactModeAuto,
},
{
name: "missing extra defaults to auto",
account: &Account{
Platform: PlatformOpenAI,
},
want: OpenAICompactModeAuto,
},
{
name: "invalid mode falls back to auto",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_mode": " invalid "},
},
want: OpenAICompactModeAuto,
},
{
name: "force on is normalized",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_mode": " FORCE_ON "},
},
want: OpenAICompactModeForceOn,
},
{
name: "force off is normalized",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_mode": "force_off"},
},
want: OpenAICompactModeForceOff,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.account.GetOpenAICompactMode(); got != tt.want {
t.Fatalf("GetOpenAICompactMode() = %q, want %q", got, tt.want)
}
})
}
}
func TestAccountOpenAICompactSupportKnown(t *testing.T) {
tests := []struct {
name string
account *Account
wantSupported bool
wantKnown bool
}{
{
name: "nil account is unknown",
wantSupported: false,
wantKnown: false,
},
{
name: "non openai account is unknown",
account: &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{"openai_compact_supported": true},
},
wantSupported: false,
wantKnown: false,
},
{
name: "force on overrides probe state",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"openai_compact_mode": OpenAICompactModeForceOn,
"openai_compact_supported": false,
},
},
wantSupported: true,
wantKnown: true,
},
{
name: "force off overrides probe state",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"openai_compact_mode": OpenAICompactModeForceOff,
"openai_compact_supported": true,
},
},
wantSupported: false,
wantKnown: true,
},
{
name: "auto true is known supported",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_supported": true},
},
wantSupported: true,
wantKnown: true,
},
{
name: "auto false is known unsupported",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_supported": false},
},
wantSupported: false,
wantKnown: true,
},
{
name: "auto without probe state remains unknown",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{},
},
wantSupported: false,
wantKnown: false,
},
{
name: "invalid probe field remains unknown",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_supported": "true"},
},
wantSupported: false,
wantKnown: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotSupported, gotKnown := tt.account.OpenAICompactSupportKnown()
if gotSupported != tt.wantSupported || gotKnown != tt.wantKnown {
t.Fatalf("OpenAICompactSupportKnown() = (%v, %v), want (%v, %v)", gotSupported, gotKnown, tt.wantSupported, tt.wantKnown)
}
})
}
}
func TestAccountAllowsOpenAICompact(t *testing.T) {
tests := []struct {
name string
account *Account
want bool
}{
{
name: "nil account does not allow compact",
want: false,
},
{
name: "non openai account does not allow compact",
account: &Account{
Platform: PlatformAnthropic,
},
want: false,
},
{
name: "unknown openai account remains allowed",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{},
},
want: true,
},
{
name: "supported openai account is allowed",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_supported": true},
},
want: true,
},
{
name: "unsupported openai account is rejected",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_supported": false},
},
want: false,
},
{
name: "force on is allowed",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOn},
},
want: true,
},
{
name: "force off is rejected",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOff},
},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.account.AllowsOpenAICompact(); got != tt.want {
t.Fatalf("AllowsOpenAICompact() = %v, want %v", got, tt.want)
}
})
}
}
func TestAccountGetCompactModelMapping(t *testing.T) {
tests := []struct {
name string
account *Account
want map[string]string
}{
{
name: "nil account returns nil",
want: nil,
},
{
name: "missing credentials returns nil",
account: &Account{
Platform: PlatformOpenAI,
},
want: nil,
},
{
name: "map any is converted",
account: &Account{
Credentials: map[string]any{
"compact_model_mapping": map[string]any{
"gpt-5.4": "gpt-5.4-openai-compact",
"invalid": 1,
},
},
},
want: map[string]string{
"gpt-5.4": "gpt-5.4-openai-compact",
},
},
{
name: "map string string is copied",
account: &Account{
Credentials: map[string]any{
"compact_model_mapping": map[string]string{
"gpt-*": "compact-*",
},
},
},
want: map[string]string{
"gpt-*": "compact-*",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.account.GetCompactModelMapping()
if !equalStringMap(got, tt.want) {
t.Fatalf("GetCompactModelMapping() = %#v, want %#v", got, tt.want)
}
})
}
}
func TestAccountResolveCompactMappedModel(t *testing.T) {
tests := []struct {
name string
credentials map[string]any
requestedModel string
expectedModel string
expectedMatch bool
}{
{
name: "no compact mapping reports unmatched",
credentials: nil,
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: false,
},
{
name: "exact compact mapping matches",
credentials: map[string]any{
"compact_model_mapping": map[string]any{
"gpt-5.4": "gpt-5.4-openai-compact",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4-openai-compact",
expectedMatch: true,
},
{
name: "exact passthrough counts as match",
credentials: map[string]any{
"compact_model_mapping": map[string]any{
"gpt-5.4": "gpt-5.4",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: true,
},
{
name: "longest wildcard wins",
credentials: map[string]any{
"compact_model_mapping": map[string]any{
"gpt-*": "fallback-compact",
"gpt-5.4*": "gpt-5.4-openai-compact",
"gpt-5.4-mini*": "gpt-5.4-mini-openai-compact",
},
},
requestedModel: "gpt-5.4-mini",
expectedModel: "gpt-5.4-mini-openai-compact",
expectedMatch: true,
},
{
name: "missing compact mapping reports unmatched",
credentials: map[string]any{
"compact_model_mapping": map[string]any{
"gpt-5.3": "gpt-5.3-openai-compact",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Credentials: tt.credentials,
}
gotModel, gotMatch := account.ResolveCompactMappedModel(tt.requestedModel)
if gotModel != tt.expectedModel || gotMatch != tt.expectedMatch {
t.Fatalf("ResolveCompactMappedModel(%q) = (%q, %v), want (%q, %v)", tt.requestedModel, gotModel, gotMatch, tt.expectedModel, tt.expectedMatch)
}
})
}
}
func equalStringMap(left, right map[string]string) bool {
if len(left) != len(right) {
return false
}
for key, want := range right {
if got, ok := left[key]; !ok || got != want {
return false
}
}
return true
}
@@ -0,0 +1,326 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAccount_IsOpenAIPassthroughEnabled(t *testing.T) {
t.Run("新字段开启", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"openai_passthrough": true,
},
}
require.True(t, account.IsOpenAIPassthroughEnabled())
})
t.Run("兼容旧字段", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_passthrough": true,
},
}
require.True(t, account.IsOpenAIPassthroughEnabled())
})
t.Run("非OpenAI账号始终关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_passthrough": true,
},
}
require.False(t, account.IsOpenAIPassthroughEnabled())
})
t.Run("空额外配置默认关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
}
require.False(t, account.IsOpenAIPassthroughEnabled())
})
}
func TestAccount_IsOpenAIOAuthPassthroughEnabled(t *testing.T) {
t.Run("仅OAuth类型允许返回开启", func(t *testing.T) {
oauthAccount := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_passthrough": true,
},
}
require.True(t, oauthAccount.IsOpenAIOAuthPassthroughEnabled())
apiKeyAccount := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"openai_passthrough": true,
},
}
require.False(t, apiKeyAccount.IsOpenAIOAuthPassthroughEnabled())
})
}
func TestAccount_IsCodexCLIOnlyEnabled(t *testing.T) {
t.Run("OpenAI OAuth 开启", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_cli_only": true,
},
}
require.True(t, account.IsCodexCLIOnlyEnabled())
})
t.Run("OpenAI OAuth 关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_cli_only": false,
},
}
require.False(t, account.IsCodexCLIOnlyEnabled())
})
t.Run("字段缺失默认关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{},
}
require.False(t, account.IsCodexCLIOnlyEnabled())
})
t.Run("类型非法默认关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_cli_only": "true",
},
}
require.False(t, account.IsCodexCLIOnlyEnabled())
})
t.Run("非 OAuth 账号始终关闭", func(t *testing.T) {
apiKeyAccount := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"codex_cli_only": true,
},
}
require.False(t, apiKeyAccount.IsCodexCLIOnlyEnabled())
otherPlatform := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_cli_only": true,
},
}
require.False(t, otherPlatform.IsCodexCLIOnlyEnabled())
})
}
func TestAccount_IsOpenAIResponsesWebSocketV2Enabled(t *testing.T) {
t.Run("OAuth使用OAuth专用开关", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_enabled": true,
},
}
require.True(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
t.Run("API Key使用API Key专用开关", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
},
}
require.True(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
t.Run("OAuth账号不会读取API Key专用开关", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
},
}
require.False(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
t.Run("分类型新键优先于兼容键", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_enabled": false,
"responses_websockets_v2_enabled": true,
"openai_ws_enabled": true,
},
}
require.False(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
t.Run("分类型键缺失时回退兼容键", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"responses_websockets_v2_enabled": true,
},
}
require.True(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
t.Run("非OpenAI账号默认关闭", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"responses_websockets_v2_enabled": true,
},
}
require.False(t, account.IsOpenAIResponsesWebSocketV2Enabled())
})
}
func TestAccount_ResolveOpenAIResponsesWebSocketV2Mode(t *testing.T) {
t.Run("default fallback to ctx_pool", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{},
}
require.Equal(t, OpenAIWSIngressModeCtxPool, account.ResolveOpenAIResponsesWebSocketV2Mode(""))
require.Equal(t, OpenAIWSIngressModeCtxPool, account.ResolveOpenAIResponsesWebSocketV2Mode("invalid"))
})
t.Run("oauth mode field has highest priority", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
"openai_oauth_responses_websockets_v2_enabled": false,
"responses_websockets_v2_enabled": false,
},
}
require.Equal(t, OpenAIWSIngressModePassthrough, account.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeCtxPool))
})
t.Run("oauth mode supports http_bridge", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModeHTTPBridge,
},
}
require.Equal(t, OpenAIWSIngressModeHTTPBridge, account.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeCtxPool))
})
t.Run("legacy enabled maps to ctx_pool", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"responses_websockets_v2_enabled": true,
},
}
require.Equal(t, OpenAIWSIngressModeCtxPool, account.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeOff))
})
t.Run("shared/dedicated mode strings are compatible with ctx_pool", func(t *testing.T) {
shared := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModeShared,
},
}
dedicated := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModeDedicated,
},
}
require.Equal(t, OpenAIWSIngressModeShared, shared.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeOff))
require.Equal(t, OpenAIWSIngressModeDedicated, dedicated.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeOff))
require.Equal(t, OpenAIWSIngressModeCtxPool, normalizeOpenAIWSIngressDefaultMode(OpenAIWSIngressModeShared))
require.Equal(t, OpenAIWSIngressModeCtxPool, normalizeOpenAIWSIngressDefaultMode(OpenAIWSIngressModeDedicated))
})
t.Run("legacy disabled maps to off", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": false,
"responses_websockets_v2_enabled": true,
},
}
require.Equal(t, OpenAIWSIngressModeOff, account.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeCtxPool))
})
t.Run("non openai always off", func(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModeDedicated,
},
}
require.Equal(t, OpenAIWSIngressModeOff, account.ResolveOpenAIResponsesWebSocketV2Mode(OpenAIWSIngressModeDedicated))
})
}
func TestAccount_OpenAIWSExtraFlags(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_ws_force_http": true,
"openai_ws_allow_store_recovery": true,
},
}
require.True(t, account.IsOpenAIWSForceHTTPEnabled())
require.True(t, account.IsOpenAIWSAllowStoreRecoveryEnabled())
off := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Extra: map[string]any{}}
require.False(t, off.IsOpenAIWSForceHTTPEnabled())
require.False(t, off.IsOpenAIWSAllowStoreRecoveryEnabled())
var nilAccount *Account
require.False(t, nilAccount.IsOpenAIWSAllowStoreRecoveryEnabled())
nonOpenAI := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_ws_allow_store_recovery": true,
},
}
require.False(t, nonOpenAI.IsOpenAIWSAllowStoreRecoveryEnabled())
}
@@ -0,0 +1,117 @@
//go:build unit
package service
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestGetPoolModeRetryCount(t *testing.T) {
tests := []struct {
name string
account *Account
expected int
}{
{
name: "default_when_not_pool_mode",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{},
},
expected: defaultPoolModeRetryCount,
},
{
name: "default_when_missing_retry_count",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
},
},
expected: defaultPoolModeRetryCount,
},
{
name: "supports_float64_from_json_credentials",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": float64(5),
},
},
expected: 5,
},
{
name: "supports_json_number",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": json.Number("4"),
},
},
expected: 4,
},
{
name: "supports_string_value",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": "2",
},
},
expected: 2,
},
{
name: "negative_value_is_clamped_to_zero",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": -1,
},
},
expected: 0,
},
{
name: "oversized_value_is_clamped_to_max",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": 99,
},
},
expected: maxPoolModeRetryCount,
},
{
name: "invalid_value_falls_back_to_default",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": "oops",
},
},
expected: defaultPoolModeRetryCount,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetPoolModeRetryCount())
})
}
}
@@ -0,0 +1,193 @@
//go:build unit
package service
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestGetPoolModeRetryStatusCodes(t *testing.T) {
tests := []struct {
name string
account *Account
expected []int
}{
{
name: "nil_account_returns_nil",
account: nil,
expected: nil,
},
{
name: "nil_credentials_returns_nil",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
},
expected: nil,
},
{
name: "missing_key_returns_nil",
account: &Account{
Type: AccountTypeAPIKey,
Platform: PlatformOpenAI,
Credentials: map[string]any{"pool_mode": true},
},
expected: nil,
},
{
name: "empty_slice_is_preserved",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{},
},
},
expected: []int{},
},
{
name: "float64_values_from_json_are_normalized",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{float64(429), float64(401), float64(403)},
},
},
expected: []int{401, 403, 429},
},
{
name: "json_number_values_supported",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{json.Number("502"), json.Number("503")},
},
},
expected: []int{502, 503},
},
{
name: "string_values_supported",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{"520", "529"},
},
},
expected: []int{520, 529},
},
{
name: "duplicates_are_deduped",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{float64(429), float64(429), float64(401)},
},
},
expected: []int{401, 429},
},
{
name: "out_of_range_values_dropped",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{float64(99), float64(600), float64(429)},
},
},
expected: []int{429},
},
{
name: "invalid_string_dropped",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{"oops", float64(429)},
},
},
expected: []int{429},
},
{
name: "non_array_value_returns_nil",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": "not-an-array",
},
},
expected: nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetPoolModeRetryStatusCodes())
})
}
}
func TestIsPoolModeRetryableStatus_Account(t *testing.T) {
tests := []struct {
name string
account *Account
statusCode int
expected bool
}{
{
name: "nil_account_falls_back_to_default_401",
account: nil,
statusCode: 401,
expected: true,
},
{
name: "nil_account_falls_back_to_default_500",
account: nil,
statusCode: 500,
expected: false,
},
{
name: "unconfigured_uses_default_403",
account: &Account{
Credentials: map[string]any{"pool_mode": true},
},
statusCode: 403,
expected: true,
},
{
name: "unconfigured_uses_default_502_false",
account: &Account{
Credentials: map[string]any{"pool_mode": true},
},
statusCode: 502,
expected: false,
},
{
name: "configured_list_overrides_default_401_dropped",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{float64(502), float64(503)},
},
},
statusCode: 401,
expected: false,
},
{
name: "configured_list_overrides_default_502_added",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{float64(502), float64(503)},
},
},
statusCode: 502,
expected: true,
},
{
name: "empty_list_disables_all_default_codes",
account: &Account{
Credentials: map[string]any{
"pool_mode_retry_status_codes": []any{},
},
},
statusCode: 429,
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.IsPoolModeRetryableStatus(tt.statusCode))
})
}
}
@@ -0,0 +1,567 @@
//go:build unit
package service
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// nextFixedDailyReset
// ---------------------------------------------------------------------------
func TestNextFixedDailyReset_BeforeResetHour(t *testing.T) {
tz := time.UTC
// 2026-03-14 06:00 UTC, reset hour = 9
after := time.Date(2026, 3, 14, 6, 0, 0, 0, tz)
got := nextFixedDailyReset(9, tz, after)
want := time.Date(2026, 3, 14, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedDailyReset_AtResetHour(t *testing.T) {
tz := time.UTC
// Exactly at reset hour → should return tomorrow
after := time.Date(2026, 3, 14, 9, 0, 0, 0, tz)
got := nextFixedDailyReset(9, tz, after)
want := time.Date(2026, 3, 15, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedDailyReset_AfterResetHour(t *testing.T) {
tz := time.UTC
// After reset hour → should return tomorrow
after := time.Date(2026, 3, 14, 15, 30, 0, 0, tz)
got := nextFixedDailyReset(9, tz, after)
want := time.Date(2026, 3, 15, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedDailyReset_MidnightReset(t *testing.T) {
tz := time.UTC
// Reset at hour 0 (midnight), currently 23:59
after := time.Date(2026, 3, 14, 23, 59, 0, 0, tz)
got := nextFixedDailyReset(0, tz, after)
want := time.Date(2026, 3, 15, 0, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedDailyReset_NonUTCTimezone(t *testing.T) {
tz, err := time.LoadLocation("Asia/Shanghai")
require.NoError(t, err)
// 2026-03-14 07:00 UTC = 2026-03-14 15:00 CST, reset hour = 9 (CST)
after := time.Date(2026, 3, 14, 7, 0, 0, 0, time.UTC)
got := nextFixedDailyReset(9, tz, after)
// Already past 9:00 CST today → tomorrow 9:00 CST = 2026-03-15 01:00 UTC
want := time.Date(2026, 3, 15, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
// ---------------------------------------------------------------------------
// lastFixedDailyReset
// ---------------------------------------------------------------------------
func TestLastFixedDailyReset_BeforeResetHour(t *testing.T) {
tz := time.UTC
now := time.Date(2026, 3, 14, 6, 0, 0, 0, tz)
got := lastFixedDailyReset(9, tz, now)
// Before today's 9:00 → yesterday 9:00
want := time.Date(2026, 3, 13, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestLastFixedDailyReset_AtResetHour(t *testing.T) {
tz := time.UTC
now := time.Date(2026, 3, 14, 9, 0, 0, 0, tz)
got := lastFixedDailyReset(9, tz, now)
// At exactly 9:00 → today 9:00
want := time.Date(2026, 3, 14, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestLastFixedDailyReset_AfterResetHour(t *testing.T) {
tz := time.UTC
now := time.Date(2026, 3, 14, 15, 0, 0, 0, tz)
got := lastFixedDailyReset(9, tz, now)
// After 9:00 → today 9:00
want := time.Date(2026, 3, 14, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
// ---------------------------------------------------------------------------
// nextFixedWeeklyReset
// ---------------------------------------------------------------------------
func TestNextFixedWeeklyReset_TargetDayAhead(t *testing.T) {
tz := time.UTC
// 2026-03-14 is Saturday (day=6), target = Monday (day=1), hour = 9
after := time.Date(2026, 3, 14, 10, 0, 0, 0, tz)
got := nextFixedWeeklyReset(1, 9, tz, after)
// Next Monday = 2026-03-16
want := time.Date(2026, 3, 16, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedWeeklyReset_TargetDayToday_BeforeHour(t *testing.T) {
tz := time.UTC
// 2026-03-16 is Monday (day=1), target = Monday, hour = 9, before 9:00
after := time.Date(2026, 3, 16, 6, 0, 0, 0, tz)
got := nextFixedWeeklyReset(1, 9, tz, after)
// Today at 9:00
want := time.Date(2026, 3, 16, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedWeeklyReset_TargetDayToday_AtHour(t *testing.T) {
tz := time.UTC
// 2026-03-16 is Monday, target = Monday, hour = 9, exactly at 9:00
after := time.Date(2026, 3, 16, 9, 0, 0, 0, tz)
got := nextFixedWeeklyReset(1, 9, tz, after)
// Next Monday at 9:00
want := time.Date(2026, 3, 23, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedWeeklyReset_TargetDayToday_AfterHour(t *testing.T) {
tz := time.UTC
// 2026-03-16 is Monday, target = Monday, hour = 9, after 9:00
after := time.Date(2026, 3, 16, 15, 0, 0, 0, tz)
got := nextFixedWeeklyReset(1, 9, tz, after)
// Next Monday at 9:00
want := time.Date(2026, 3, 23, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedWeeklyReset_TargetDayPast(t *testing.T) {
tz := time.UTC
// 2026-03-18 is Wednesday (day=3), target = Monday (day=1)
after := time.Date(2026, 3, 18, 10, 0, 0, 0, tz)
got := nextFixedWeeklyReset(1, 9, tz, after)
// Next Monday = 2026-03-23
want := time.Date(2026, 3, 23, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestNextFixedWeeklyReset_Sunday(t *testing.T) {
tz := time.UTC
// 2026-03-14 is Saturday (day=6), target = Sunday (day=0)
after := time.Date(2026, 3, 14, 10, 0, 0, 0, tz)
got := nextFixedWeeklyReset(0, 0, tz, after)
// Next Sunday = 2026-03-15
want := time.Date(2026, 3, 15, 0, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
// ---------------------------------------------------------------------------
// lastFixedWeeklyReset
// ---------------------------------------------------------------------------
func TestLastFixedWeeklyReset_SameDay_AfterHour(t *testing.T) {
tz := time.UTC
// 2026-03-16 is Monday (day=1), target = Monday, hour = 9, now = 15:00
now := time.Date(2026, 3, 16, 15, 0, 0, 0, tz)
got := lastFixedWeeklyReset(1, 9, tz, now)
// Today at 9:00
want := time.Date(2026, 3, 16, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestLastFixedWeeklyReset_SameDay_BeforeHour(t *testing.T) {
tz := time.UTC
// 2026-03-16 is Monday, target = Monday, hour = 9, now = 06:00
now := time.Date(2026, 3, 16, 6, 0, 0, 0, tz)
got := lastFixedWeeklyReset(1, 9, tz, now)
// Last Monday at 9:00 = 2026-03-09
want := time.Date(2026, 3, 9, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
func TestLastFixedWeeklyReset_DifferentDay(t *testing.T) {
tz := time.UTC
// 2026-03-18 is Wednesday (day=3), target = Monday (day=1)
now := time.Date(2026, 3, 18, 10, 0, 0, 0, tz)
got := lastFixedWeeklyReset(1, 9, tz, now)
// Last Monday = 2026-03-16
want := time.Date(2026, 3, 16, 9, 0, 0, 0, tz)
assert.Equal(t, want, got)
}
// ---------------------------------------------------------------------------
// isFixedDailyPeriodExpired
// ---------------------------------------------------------------------------
func TestIsFixedDailyPeriodExpired_ZeroPeriodStart(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
assert.True(t, a.isFixedDailyPeriodExpired(time.Time{}))
}
func TestIsFixedDailyPeriodExpired_NotExpired(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
// Anchor periodStart to today's 12:00 UTC: always strictly after today's
// 09:00 UTC reset (and yesterday's). Using time.Now().Add(-1*time.Minute)
// is flaky inside the 09:00-09:01 UTC reset window.
now := time.Now().UTC()
periodStart := time.Date(now.Year(), now.Month(), now.Day(), 12, 0, 0, 0, time.UTC)
assert.False(t, a.isFixedDailyPeriodExpired(periodStart))
}
func TestIsFixedDailyPeriodExpired_Expired(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
// Period started 3 days ago → definitely expired
periodStart := time.Now().Add(-72 * time.Hour)
assert.True(t, a.isFixedDailyPeriodExpired(periodStart))
}
func TestIsFixedDailyPeriodExpired_InvalidTimezone(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "Invalid/Timezone",
}}
// Invalid timezone falls back to UTC
periodStart := time.Now().Add(-72 * time.Hour)
assert.True(t, a.isFixedDailyPeriodExpired(periodStart))
}
// ---------------------------------------------------------------------------
// isFixedWeeklyPeriodExpired
// ---------------------------------------------------------------------------
func TestIsFixedWeeklyPeriodExpired_ZeroPeriodStart(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
assert.True(t, a.isFixedWeeklyPeriodExpired(time.Time{}))
}
func TestIsFixedWeeklyPeriodExpired_NotExpired(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
// Anchor periodStart to today's 12:00 UTC: always strictly after the most
// recent Monday 09:00 UTC reset, regardless of which weekday/hour the test
// runs. Using time.Now().Add(-1*time.Minute) is flaky inside the
// Monday 09:00-09:01 UTC reset window.
now := time.Now().UTC()
periodStart := time.Date(now.Year(), now.Month(), now.Day(), 12, 0, 0, 0, time.UTC)
assert.False(t, a.isFixedWeeklyPeriodExpired(periodStart))
}
func TestIsFixedWeeklyPeriodExpired_Expired(t *testing.T) {
a := &Account{Extra: map[string]any{
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}}
// Period started 10 days ago → definitely expired
periodStart := time.Now().Add(-240 * time.Hour)
assert.True(t, a.isFixedWeeklyPeriodExpired(periodStart))
}
// ---------------------------------------------------------------------------
// ValidateQuotaResetConfig
// ---------------------------------------------------------------------------
func TestValidateQuotaResetConfig_NilExtra(t *testing.T) {
assert.NoError(t, ValidateQuotaResetConfig(nil))
}
func TestValidateQuotaResetConfig_EmptyExtra(t *testing.T) {
assert.NoError(t, ValidateQuotaResetConfig(map[string]any{}))
}
func TestValidateQuotaResetConfig_ValidFixed(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "Asia/Shanghai",
}
assert.NoError(t, ValidateQuotaResetConfig(extra))
}
func TestValidateQuotaResetConfig_ValidRolling(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "rolling",
"quota_weekly_reset_mode": "rolling",
}
assert.NoError(t, ValidateQuotaResetConfig(extra))
}
func TestValidateQuotaResetConfig_InvalidTimezone(t *testing.T) {
extra := map[string]any{
"quota_reset_timezone": "Not/A/Timezone",
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_reset_timezone")
}
func TestValidateQuotaResetConfig_InvalidDailyMode(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "invalid",
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_daily_reset_mode")
}
func TestValidateQuotaResetConfig_InvalidDailyHour_TooHigh(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_hour": float64(24),
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_daily_reset_hour")
}
func TestValidateQuotaResetConfig_InvalidDailyHour_Negative(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_hour": float64(-1),
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_daily_reset_hour")
}
func TestValidateQuotaResetConfig_InvalidWeeklyMode(t *testing.T) {
extra := map[string]any{
"quota_weekly_reset_mode": "unknown",
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_weekly_reset_mode")
}
func TestValidateQuotaResetConfig_InvalidWeeklyDay_TooHigh(t *testing.T) {
extra := map[string]any{
"quota_weekly_reset_day": float64(7),
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_weekly_reset_day")
}
func TestValidateQuotaResetConfig_InvalidWeeklyDay_Negative(t *testing.T) {
extra := map[string]any{
"quota_weekly_reset_day": float64(-1),
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_weekly_reset_day")
}
func TestValidateQuotaResetConfig_InvalidWeeklyHour(t *testing.T) {
extra := map[string]any{
"quota_weekly_reset_hour": float64(25),
}
err := ValidateQuotaResetConfig(extra)
require.Error(t, err)
assert.Contains(t, err.Error(), "quota_weekly_reset_hour")
}
func TestValidateQuotaResetConfig_BoundaryValues(t *testing.T) {
// All boundary values should be valid
extra := map[string]any{
"quota_daily_reset_hour": float64(23),
"quota_weekly_reset_day": float64(0), // Sunday
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "UTC",
}
assert.NoError(t, ValidateQuotaResetConfig(extra))
extra2 := map[string]any{
"quota_daily_reset_hour": float64(0),
"quota_weekly_reset_day": float64(6), // Saturday
"quota_weekly_reset_hour": float64(23),
}
assert.NoError(t, ValidateQuotaResetConfig(extra2))
}
// ---------------------------------------------------------------------------
// NormalizeFixedQuotaWindows
// ---------------------------------------------------------------------------
func TestNormalizeFixedQuotaWindows_ClearsExpiredWeeklyWindow(t *testing.T) {
now := time.Now().UTC()
daysSinceMonday := (int(now.Weekday()) + 6) % 7
currentWeekStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, -daysSinceMonday)
staleStart := currentWeekStart.Add(-24 * time.Hour)
extra := map[string]any{
"quota_weekly_limit": 500.0,
"quota_weekly_used": 76.0,
"quota_weekly_start": staleStart.Format(time.RFC3339),
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "UTC",
}
NormalizeFixedQuotaWindows(extra)
assert.Equal(t, 0.0, extra["quota_weekly_used"])
assert.Equal(t, currentWeekStart.Format(time.RFC3339), extra["quota_weekly_start"])
}
func TestNormalizeFixedQuotaWindows_KeepsActiveWeeklyWindow(t *testing.T) {
now := time.Now().UTC()
daysSinceMonday := (int(now.Weekday()) + 6) % 7
currentWeekStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, -daysSinceMonday)
extra := map[string]any{
"quota_weekly_limit": 500.0,
"quota_weekly_used": 76.0,
"quota_weekly_start": currentWeekStart.Format(time.RFC3339),
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "UTC",
}
NormalizeFixedQuotaWindows(extra)
assert.Equal(t, 76.0, extra["quota_weekly_used"])
assert.Equal(t, currentWeekStart.Format(time.RFC3339), extra["quota_weekly_start"])
}
// ---------------------------------------------------------------------------
// ComputeQuotaResetAt
// ---------------------------------------------------------------------------
func TestComputeQuotaResetAt_RollingMode_NoResetAt(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "rolling",
"quota_weekly_reset_mode": "rolling",
}
ComputeQuotaResetAt(extra)
_, hasDailyResetAt := extra["quota_daily_reset_at"]
_, hasWeeklyResetAt := extra["quota_weekly_reset_at"]
assert.False(t, hasDailyResetAt, "rolling mode should not set quota_daily_reset_at")
assert.False(t, hasWeeklyResetAt, "rolling mode should not set quota_weekly_reset_at")
}
func TestComputeQuotaResetAt_RollingMode_ClearsExistingResetAt(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "rolling",
"quota_weekly_reset_mode": "rolling",
"quota_daily_reset_at": "2026-03-14T09:00:00Z",
"quota_weekly_reset_at": "2026-03-16T09:00:00Z",
}
ComputeQuotaResetAt(extra)
_, hasDailyResetAt := extra["quota_daily_reset_at"]
_, hasWeeklyResetAt := extra["quota_weekly_reset_at"]
assert.False(t, hasDailyResetAt, "rolling mode should remove quota_daily_reset_at")
assert.False(t, hasWeeklyResetAt, "rolling mode should remove quota_weekly_reset_at")
}
func TestComputeQuotaResetAt_FixedDaily_SetsResetAt(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "UTC",
}
ComputeQuotaResetAt(extra)
resetAtStr, ok := extra["quota_daily_reset_at"].(string)
require.True(t, ok, "quota_daily_reset_at should be set")
resetAt, err := time.Parse(time.RFC3339, resetAtStr)
require.NoError(t, err)
// Reset time should be in the future
assert.True(t, resetAt.After(time.Now()), "reset_at should be in the future")
// Reset hour should be 9 UTC
assert.Equal(t, 9, resetAt.UTC().Hour())
}
func TestComputeQuotaResetAt_FixedWeekly_SetsResetAt(t *testing.T) {
extra := map[string]any{
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1), // Monday
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "UTC",
}
ComputeQuotaResetAt(extra)
resetAtStr, ok := extra["quota_weekly_reset_at"].(string)
require.True(t, ok, "quota_weekly_reset_at should be set")
resetAt, err := time.Parse(time.RFC3339, resetAtStr)
require.NoError(t, err)
// Reset time should be in the future
assert.True(t, resetAt.After(time.Now()), "reset_at should be in the future")
// Reset day should be Monday
assert.Equal(t, time.Monday, resetAt.UTC().Weekday())
}
func TestComputeQuotaResetAt_FixedDaily_WithTimezone(t *testing.T) {
tz, err := time.LoadLocation("Asia/Shanghai")
require.NoError(t, err)
extra := map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(9),
"quota_reset_timezone": "Asia/Shanghai",
}
ComputeQuotaResetAt(extra)
resetAtStr, ok := extra["quota_daily_reset_at"].(string)
require.True(t, ok)
resetAt, err := time.Parse(time.RFC3339, resetAtStr)
require.NoError(t, err)
// In Shanghai timezone, the hour should be 9
assert.Equal(t, 9, resetAt.In(tz).Hour())
}
func TestComputeQuotaResetAt_DefaultTimezone(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(12),
}
ComputeQuotaResetAt(extra)
resetAtStr, ok := extra["quota_daily_reset_at"].(string)
require.True(t, ok)
resetAt, err := time.Parse(time.RFC3339, resetAtStr)
require.NoError(t, err)
// Default timezone is UTC
assert.Equal(t, 12, resetAt.UTC().Hour())
}
func TestComputeQuotaResetAt_InvalidHour_ClampedToZero(t *testing.T) {
extra := map[string]any{
"quota_daily_reset_mode": "fixed",
"quota_daily_reset_hour": float64(99),
"quota_reset_timezone": "UTC",
}
ComputeQuotaResetAt(extra)
resetAtStr, ok := extra["quota_daily_reset_at"].(string)
require.True(t, ok)
resetAt, err := time.Parse(time.RFC3339, resetAtStr)
require.NoError(t, err)
// Invalid hour → clamped to 0
assert.Equal(t, 0, resetAt.UTC().Hour())
}
@@ -0,0 +1,123 @@
//go:build unit
package service
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestAccountIsSchedulable_QuotaExceeded(t *testing.T) {
now := time.Now()
tests := []struct {
name string
account *Account
want bool
}{
{
name: "apikey daily quota exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"quota_daily_limit": 10.0,
"quota_daily_used": 10.0,
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
},
},
want: false,
},
{
name: "apikey weekly quota exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"quota_weekly_limit": 50.0,
"quota_weekly_used": 50.0,
"quota_weekly_start": now.Add(-2 * 24 * time.Hour).Format(time.RFC3339),
},
},
want: false,
},
{
name: "apikey total quota exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"quota_limit": 100.0,
"quota_used": 100.0,
},
},
want: false,
},
{
name: "apikey quota not exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"quota_daily_limit": 10.0,
"quota_daily_used": 5.0,
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
},
},
want: true,
},
{
name: "apikey expired daily period restores schedulable",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeAPIKey,
Extra: map[string]any{
"quota_daily_limit": 10.0,
"quota_daily_used": 10.0,
"quota_daily_start": now.Add(-25 * time.Hour).Format(time.RFC3339),
},
},
want: true,
},
{
name: "oauth ignores quota exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeOAuth,
Extra: map[string]any{
"quota_daily_limit": 10.0,
"quota_daily_used": 10.0,
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
},
},
want: true,
},
{
name: "bedrock quota exceeded",
account: &Account{
Status: StatusActive,
Schedulable: true,
Type: AccountTypeBedrock,
Extra: map[string]any{
"quota_limit": 200.0,
"quota_used": 200.0,
},
},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, tt.account.IsSchedulable())
})
}
}
@@ -0,0 +1,25 @@
//go:build unit
package service
import "context"
func (s *accountRepoStub) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (r *openAIAccountTestRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *groupAwareMockAccountRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *mockAccountRepoForPlatform) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *mockAccountRepoForGemini) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
return m.ListActive(context.Background())
}
@@ -0,0 +1,139 @@
package service
import (
"encoding/json"
"testing"
)
func TestGetBaseRPM(t *testing.T) {
tests := []struct {
name string
extra map[string]any
expected int
}{
{"nil extra", nil, 0},
{"no key", map[string]any{}, 0},
{"zero", map[string]any{"base_rpm": 0}, 0},
{"int value", map[string]any{"base_rpm": 15}, 15},
{"float value", map[string]any{"base_rpm": 15.0}, 15},
{"string value", map[string]any{"base_rpm": "15"}, 15},
{"negative value", map[string]any{"base_rpm": -5}, 0},
{"int64 value", map[string]any{"base_rpm": int64(20)}, 20},
{"json.Number value", map[string]any{"base_rpm": json.Number("25")}, 25},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := &Account{Extra: tt.extra}
if got := a.GetBaseRPM(); got != tt.expected {
t.Errorf("GetBaseRPM() = %d, want %d", got, tt.expected)
}
})
}
}
func TestGetRPMStrategy(t *testing.T) {
tests := []struct {
name string
extra map[string]any
expected string
}{
{"nil extra", nil, "tiered"},
{"no key", map[string]any{}, "tiered"},
{"tiered", map[string]any{"rpm_strategy": "tiered"}, "tiered"},
{"sticky_exempt", map[string]any{"rpm_strategy": "sticky_exempt"}, "sticky_exempt"},
{"invalid", map[string]any{"rpm_strategy": "foobar"}, "tiered"},
{"empty string fallback", map[string]any{"rpm_strategy": ""}, "tiered"},
{"numeric value fallback", map[string]any{"rpm_strategy": 123}, "tiered"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := &Account{Extra: tt.extra}
if got := a.GetRPMStrategy(); got != tt.expected {
t.Errorf("GetRPMStrategy() = %q, want %q", got, tt.expected)
}
})
}
}
func TestCheckRPMSchedulability(t *testing.T) {
tests := []struct {
name string
extra map[string]any
currentRPM int
expected WindowCostSchedulability
}{
{"disabled", map[string]any{}, 100, WindowCostSchedulable},
{"green zone", map[string]any{"base_rpm": 15}, 10, WindowCostSchedulable},
{"yellow zone tiered", map[string]any{"base_rpm": 15}, 15, WindowCostStickyOnly},
{"red zone tiered", map[string]any{"base_rpm": 15}, 18, WindowCostNotSchedulable},
{"sticky_exempt at limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 15, WindowCostStickyOnly},
{"sticky_exempt over limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 100, WindowCostStickyOnly},
{"custom buffer", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 14, WindowCostStickyOnly},
{"custom buffer red", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 15, WindowCostNotSchedulable},
{"base_rpm=1 green", map[string]any{"base_rpm": 1}, 0, WindowCostSchedulable},
{"base_rpm=1 yellow (at limit)", map[string]any{"base_rpm": 1}, 1, WindowCostStickyOnly},
{"base_rpm=1 red (at limit+buffer)", map[string]any{"base_rpm": 1}, 2, WindowCostNotSchedulable},
{"negative currentRPM", map[string]any{"base_rpm": 15}, -1, WindowCostSchedulable},
{"base_rpm negative disabled", map[string]any{"base_rpm": -5}, 10, WindowCostSchedulable},
{"very high currentRPM", map[string]any{"base_rpm": 10}, 9999, WindowCostNotSchedulable},
{"sticky_exempt very high currentRPM", map[string]any{"base_rpm": 10, "rpm_strategy": "sticky_exempt"}, 9999, WindowCostStickyOnly},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := &Account{Extra: tt.extra}
if got := a.CheckRPMSchedulability(tt.currentRPM); got != tt.expected {
t.Errorf("CheckRPMSchedulability(%d) = %d, want %d", tt.currentRPM, got, tt.expected)
}
})
}
}
func TestGetRPMStickyBuffer(t *testing.T) {
tests := []struct {
name string
concurrency int
extra map[string]any
expected int
}{
// 基础退化
{"nil extra", 0, nil, 0},
{"no keys", 0, map[string]any{}, 0},
{"base_rpm=0", 0, map[string]any{"base_rpm": 0}, 0},
// 新公式: concurrency + maxSessions, floor = base/5
{"conc=3 sess=10 → 13", 3, map[string]any{"base_rpm": 15, "max_sessions": 10}, 13},
{"conc=2 sess=5 → 7", 2, map[string]any{"base_rpm": 10, "max_sessions": 5}, 7},
{"conc=3 sess=15 → 18", 3, map[string]any{"base_rpm": 30, "max_sessions": 15}, 18},
// floor 生效 (conc+sess < base/5)
{"conc=0 sess=0 base=15 → floor 3", 0, map[string]any{"base_rpm": 15}, 3},
{"conc=0 sess=0 base=10 → floor 2", 0, map[string]any{"base_rpm": 10}, 2},
{"conc=0 sess=0 base=1 → floor 1", 0, map[string]any{"base_rpm": 1}, 1},
{"conc=0 sess=0 base=4 → floor 1", 0, map[string]any{"base_rpm": 4}, 1},
{"conc=1 sess=0 base=15 → floor 3", 1, map[string]any{"base_rpm": 15}, 3},
// 手动 override
{"custom buffer=5", 3, map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5, "max_sessions": 10}, 5},
{"custom buffer=0 fallback", 3, map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 0, "max_sessions": 10}, 13},
{"custom buffer negative fallback", 3, map[string]any{"base_rpm": 10, "rpm_sticky_buffer": -1, "max_sessions": 10}, 13},
{"custom buffer with float", 3, map[string]any{"base_rpm": 10, "rpm_sticky_buffer": float64(7)}, 7},
// 负值 clamp
{"negative concurrency clamped", -5, map[string]any{"base_rpm": 15, "max_sessions": 10}, 10},
{"negative maxSessions clamped", 3, map[string]any{"base_rpm": 15, "max_sessions": -5}, 3},
// 高并发低会话
{"conc=10 sess=5 → 15", 10, map[string]any{"base_rpm": 10, "max_sessions": 5}, 15},
// json.Number
{"json.Number base_rpm", 3, map[string]any{"base_rpm": json.Number("10"), "max_sessions": json.Number("5")}, 8},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := &Account{Concurrency: tt.concurrency, Extra: tt.extra}
if got := a.GetRPMStickyBuffer(); got != tt.expected {
t.Errorf("GetRPMStickyBuffer() = %d, want %d", got, tt.expected)
}
})
}
}
@@ -0,0 +1,514 @@
package service
import (
"encoding/json"
"math"
"strconv"
"strings"
"time"
)
// AccountSchedulingThresholdDecision captures the pure pause decision for one account.
type AccountSchedulingThresholdDecision struct {
ShouldPause bool
Platform string
Window string
Scope string
ThresholdPercent int
UsedPercent float64
Until *time.Time
}
type accountSchedulingThresholdCandidate struct {
window string
scope string
usedPercent float64
until *time.Time
}
const accountSchedulingThresholdCredentialKey = "account_scheduling_threshold"
// EvaluateAccountSchedulingThreshold evaluates whether an account should be paused
// based on the current per-platform scheduling threshold snapshot.
func EvaluateAccountSchedulingThreshold(account *Account, thresholds map[string]int, now time.Time) AccountSchedulingThresholdDecision {
decision := AccountSchedulingThresholdDecision{}
if account == nil {
return decision
}
decision.Platform = strings.ToLower(strings.TrimSpace(account.Platform))
if decision.Platform == "" {
return decision
}
if !isAllowedSchedulingThresholdPlatform(decision.Platform) {
return decision
}
threshold, ok := resolveEffectiveAccountSchedulingThreshold(account, thresholds, decision.Platform)
decision.ThresholdPercent = threshold
if !ok || threshold >= 100 {
return decision
}
var winner *accountSchedulingThresholdCandidate
switch decision.Platform {
case PlatformOpenAI:
winner = pickLatestResetSchedulingCandidate(openAIThresholdCandidates(account, now), threshold, now)
case PlatformAnthropic:
winner = pickLatestResetSchedulingCandidate(anthropicThresholdCandidates(account), threshold, now)
case PlatformGrok:
winner = pickLatestResetSchedulingCandidate(grokThresholdCandidates(account), threshold, now)
case PlatformKimi:
winner = pickLatestResetSchedulingCandidate(cnProviderThresholdCandidates(account, PlatformKimi), threshold, now)
case PlatformZhipu:
winner = pickLatestResetSchedulingCandidate(cnProviderThresholdCandidates(account, PlatformZhipu), threshold, now)
default:
return decision
}
if winner == nil {
return decision
}
decision.ShouldPause = true
decision.Window = winner.window
decision.Scope = winner.scope
decision.UsedPercent = winner.usedPercent
decision.Until = winner.until
return decision
}
func isAllowedSchedulingThresholdPlatform(platform string) bool {
for _, allowed := range AllowedSchedulingThresholdPlatforms {
if platform == allowed {
return true
}
}
return false
}
func resolveEffectiveAccountSchedulingThreshold(account *Account, thresholds map[string]int, platform string) (int, bool) {
if account != nil {
if threshold, ok := accountSchedulingThresholdOverride(account); ok {
return threshold, true
}
}
return lookupAccountSchedulingThreshold(thresholds, platform)
}
func accountSchedulingThresholdOverride(account *Account) (int, bool) {
if account == nil || len(account.Credentials) == 0 {
return 0, false
}
raw, ok := account.Credentials[accountSchedulingThresholdCredentialKey]
if !ok {
return 0, false
}
return parseAccountSchedulingThresholdValue(raw)
}
func parseAccountSchedulingThresholdValue(raw any) (int, bool) {
var value int
switch v := raw.(type) {
case int:
value = v
case int64:
value = int(v)
case float64:
value = int(math.Round(v))
case float32:
value = int(math.Round(float64(v)))
case json.Number:
parsed, err := v.Float64()
if err != nil {
return 0, false
}
value = int(math.Round(parsed))
case string:
raw := strings.TrimSpace(v)
parsed, err := strconv.Atoi(raw)
if err == nil {
value = parsed
break
}
parsedFloat, floatErr := strconv.ParseFloat(raw, 64)
if floatErr != nil {
return 0, false
}
value = int(math.Round(parsedFloat))
default:
return 0, false
}
if value < 1 || value > 100 {
return 0, false
}
return value, true
}
func lookupAccountSchedulingThreshold(thresholds map[string]int, platform string) (int, bool) {
if len(thresholds) == 0 {
return 0, false
}
value, ok := thresholds[platform]
return value, ok
}
func openAIThresholdCandidates(account *Account, now time.Time) []*accountSchedulingThresholdCandidate {
if account == nil {
return nil
}
if !openAICodexSnapshotIdentityTrusted(account) {
return nil
}
return []*accountSchedulingThresholdCandidate{
openAIThresholdCandidate(account.Extra, "5h", now),
openAIThresholdCandidate(account.Extra, "7d", now),
}
}
func openAICodexSnapshotIdentityTrusted(account *Account) bool {
if account == nil || !account.IsOpenAIOAuth() || len(account.Extra) == 0 {
return true
}
if identityValuesConflict(
firstStringValue(account.Credentials, "email"),
firstStringValue(account.Extra, "email", "email_address"),
) {
return false
}
if identityValuesConflict(
firstStringValue(account.Credentials, "chatgpt_account_id"),
firstStringValue(account.Extra, "chatgpt_account_id", "account_id"),
) {
return false
}
if identityValuesConflict(
firstStringValue(account.Credentials, "workspace_id", "chatgpt_workspace_id", "organization_id", "org_id"),
firstStringValue(account.Extra, "workspace_id", "chatgpt_workspace_id", "organization_id", "org_id"),
) {
return false
}
return true
}
func identityValuesConflict(left, right string) bool {
left = strings.TrimSpace(left)
right = strings.TrimSpace(right)
return left != "" && right != "" && !strings.EqualFold(left, right)
}
// firstStringValue returns the first non-empty string among the given map keys.
// Used by OpenAI codex snapshot identity matching for scheduling thresholds.
func firstStringValue(values map[string]any, keys ...string) string {
if len(values) == 0 {
return ""
}
for _, key := range keys {
raw, ok := values[key]
if !ok || raw == nil {
continue
}
switch typed := raw.(type) {
case string:
if v := strings.TrimSpace(typed); v != "" {
return v
}
default:
if v := strings.TrimSpace(stringValue(raw)); v != "" {
return v
}
}
}
return ""
}
func openAIThresholdCandidate(extra map[string]any, window string, now time.Time) *accountSchedulingThresholdCandidate {
if len(extra) == 0 {
return nil
}
var (
usedPercentKey string
resetAtKey string
)
switch window {
case "5h":
usedPercentKey = "codex_5h_used_percent"
resetAtKey = "codex_5h_reset_at"
case "7d":
usedPercentKey = "codex_7d_used_percent"
resetAtKey = "codex_7d_reset_at"
default:
return nil
}
usedPercent, ok := extra[usedPercentKey]
if !ok {
return nil
}
if openAIQuotaWindowReset(extra, window, now) || openAICodexSnapshotStaleForPause(extra, now) {
return nil
}
return &accountSchedulingThresholdCandidate{
window: window,
usedPercent: schedulingPercentValue(usedPercent),
until: parseSchedulingResetAt(extra[resetAtKey]),
}
}
func anthropicThresholdCandidates(account *Account) []*accountSchedulingThresholdCandidate {
if account == nil {
return nil
}
var candidates []*accountSchedulingThresholdCandidate
if usedPercent := utilizationAsPercent(account.Extra["session_window_utilization"]); usedPercent > 0 {
candidates = append(candidates, &accountSchedulingThresholdCandidate{
window: "5h",
usedPercent: usedPercent,
until: cloneTimePtr(account.SessionWindowEnd),
})
}
if usedPercent := utilizationAsPercent(account.Extra["passive_usage_7d_utilization"]); usedPercent > 0 {
candidates = append(candidates, &accountSchedulingThresholdCandidate{
window: "7d",
usedPercent: usedPercent,
until: parseSchedulingResetAt(account.Extra["passive_usage_7d_reset"]),
})
}
return candidates
}
// NOTE: Gemini / Kiro / Antigravity are intentionally NOT threshold-pausing
// platforms (see AllowedSchedulingThresholdPlatforms and the evaluator switch,
// asserted by TestEvaluateAccountSchedulingThreshold_UnsupportedPlatformsDoNotPause).
// Their former per-platform candidate readers were dead code — never reachable
// from EvaluateAccountSchedulingThreshold — and have been removed to avoid the
// false impression that configuring a threshold for them has any effect. The
// kiro_sched_* / antigravity_sched_* extras are still written purely as
// observability snapshots.
// grokThresholdCandidates uses only header-projected
// grok_sched_utilization / grok_sched_reset_at (rolling quota window, reset
// capped at ~25h when written). Official billing 7d/30d windows are not used
// for auto-pause here.
func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCandidate {
if account == nil {
return nil
}
return []*accountSchedulingThresholdCandidate{
{
window: "quota",
scope: "grok",
usedPercent: schedulingPercentValue(account.Extra["grok_sched_utilization"]),
until: parseSchedulingResetAt(account.Extra["grok_sched_reset_at"]),
},
}
}
// cnProviderThresholdCandidates 读取国产供应商 Coding Plan 账号的 5h / weekly 滚动窗口
// 用量快照(由 CNProviderQuotaService 写入 account.Extra,键形如
// <provider>_5h_used_percent / <provider>_weekly_reset_at)。payg 账号无此快照,
// 候选为空 → 不触发阈值停调(余额型走余额检测)。与 openai 的快照驱动停调一致:
// 仅当用量超阈值且窗口尚未重置时才停调。
func cnProviderThresholdCandidates(account *Account, provider string) []*accountSchedulingThresholdCandidate {
if account == nil || len(account.Extra) == 0 {
return nil
}
return []*accountSchedulingThresholdCandidate{
cnThresholdCandidate(account.Extra, provider, "5h"),
cnThresholdCandidate(account.Extra, provider, "weekly"),
}
}
func cnThresholdCandidate(extra map[string]any, provider, window string) *accountSchedulingThresholdCandidate {
var usedKey, resetKey string
switch window {
case "5h":
usedKey = cnExtraKey(provider, cnExtraSuffix5hUsed)
resetKey = cnExtraKey(provider, cnExtraSuffix5hReset)
case "weekly":
usedKey = cnExtraKey(provider, cnExtraSuffixWeeklyUsed)
resetKey = cnExtraKey(provider, cnExtraSuffixWeeklyReset)
default:
return nil
}
usedPercent, ok := extra[usedKey]
if !ok {
return nil
}
return &accountSchedulingThresholdCandidate{
window: window,
scope: provider,
usedPercent: schedulingPercentValue(usedPercent),
until: parseSchedulingResetAt(extra[resetKey]),
}
}
func pickLatestResetSchedulingCandidate(candidates []*accountSchedulingThresholdCandidate, threshold int, now time.Time) *accountSchedulingThresholdCandidate {
var winner *accountSchedulingThresholdCandidate
for _, candidate := range candidates {
if !candidateMatchesThreshold(candidate, threshold, now) {
continue
}
if winner == nil || candidate.until.After(*winner.until) {
winner = candidate
continue
}
if winner.until.Equal(*candidate.until) && candidate.usedPercent > winner.usedPercent {
winner = candidate
}
}
return winner
}
func candidateMatchesThreshold(candidate *accountSchedulingThresholdCandidate, threshold int, now time.Time) bool {
if candidate == nil || candidate.until == nil || !candidate.until.After(now) {
return false
}
return candidate.usedPercent >= float64(threshold)
}
func utilizationAsPercent(raw any) float64 {
switch v := raw.(type) {
case float64:
if v >= 0 && v <= 1 {
return v * 100
}
return v
case float32:
value := float64(v)
if value >= 0 && value <= 1 {
return value * 100
}
return value
case int:
return float64(v)
case int64:
return float64(v)
case json.Number:
value, err := v.Float64()
if err != nil {
return 0
}
if strings.Contains(v.String(), ".") && value >= 0 && value <= 1 {
return value * 100
}
return value
case string:
trimmed := strings.TrimSpace(v)
value, err := strconv.ParseFloat(trimmed, 64)
if err != nil {
return 0
}
if strings.Contains(trimmed, ".") && value >= 0 && value <= 1 {
return value * 100
}
return value
default:
return 0
}
}
func schedulingPercentValue(raw any) float64 {
switch v := raw.(type) {
case float64:
return v
case float32:
return float64(v)
case int:
return float64(v)
case int64:
return float64(v)
case json.Number:
value, err := v.Float64()
if err != nil {
return 0
}
return value
case string:
value, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
if err != nil {
return 0
}
return value
default:
return 0
}
}
func parseSchedulingResetAt(raw any) *time.Time {
switch v := raw.(type) {
case nil:
return nil
case time.Time:
ts := v
return &ts
case *time.Time:
return cloneTimePtr(v)
case string:
trimmed := strings.TrimSpace(v)
if trimmed == "" {
return nil
}
ts, err := parseSchedulingTime(trimmed)
if err != nil {
return nil
}
return &ts
case json.Number:
if value, err := v.Int64(); err == nil && value > 0 {
ts := time.Unix(value, 0)
return &ts
}
if value, err := v.Float64(); err == nil && value > 0 {
ts := time.Unix(int64(value), 0)
return &ts
}
case float64:
if v > 0 {
ts := time.Unix(int64(v), 0)
return &ts
}
case float32:
if v > 0 {
ts := time.Unix(int64(v), 0)
return &ts
}
case int:
if v > 0 {
ts := time.Unix(int64(v), 0)
return &ts
}
case int64:
if v > 0 {
ts := time.Unix(v, 0)
return &ts
}
}
return nil
}
func parseSchedulingTime(raw string) (time.Time, error) {
formats := []string{
time.RFC3339,
time.RFC3339Nano,
"2006-01-02T15:04:05Z",
"2006-01-02T15:04:05.000Z",
}
for _, format := range formats {
if ts, err := time.Parse(format, raw); err == nil {
return ts, nil
}
}
return time.Time{}, strconv.ErrSyntax
}
func cloneTimePtr(src *time.Time) *time.Time {
if src == nil {
return nil
}
value := *src
return &value
}
@@ -0,0 +1,436 @@
//go:build unit
package service
import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
func TestEvaluateAccountSchedulingThreshold_OpenAIChoosesLatestResetWindow(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
wantUntil := now.Add(72 * time.Hour)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_5h_used_percent": 90.0,
"codex_5h_reset_at": now.Add(2 * time.Hour).Format(time.RFC3339),
"codex_7d_used_percent": 85.0,
"codex_7d_reset_at": wantUntil.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformOpenAI: 80,
}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, PlatformOpenAI, decision.Platform)
require.Equal(t, "7d", decision.Window)
require.Empty(t, decision.Scope)
require.Equal(t, 85.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, wantUntil.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_OpenAIIgnoresMismatchedCodexSnapshotIdentity(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 13, 8, 50, 0, 0, time.UTC)
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"email": "CageLeen9208@outlook.com",
"chatgpt_account_id": "1f945aa7-d9a9-4369-9542-0c702ff4adb0",
"workspace_id": "org-nU4goUxMmureroyswT5oYPv4",
"chatgpt_workspace_id": "org-nU4goUxMmureroyswT5oYPv4",
},
Extra: map[string]any{
"email": "MasonDobies01@outlook.com",
"name": "Paul Clark",
"workspace_id": "org-avRk1G4qdXg7qph3cRIraNKf",
"codex_7d_used_percent": 100.0,
"codex_7d_reset_at": now.Add(7 * 24 * time.Hour).Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformOpenAI: 99,
}, now)
require.False(t, decision.ShouldPause)
}
func TestEvaluateAccountSchedulingThreshold_AnthropicIgnoresExpiredFiveHourWindow(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
expiredEnd := now.Add(-30 * time.Minute)
wantUntil := now.Add(5 * 24 * time.Hour)
account := &Account{
Platform: PlatformAnthropic,
SessionWindowEnd: &expiredEnd,
Extra: map[string]any{
"session_window_utilization": 0.99,
"passive_usage_7d_utilization": 0.82,
"passive_usage_7d_reset": float64(wantUntil.Unix()),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformAnthropic: 80,
}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, PlatformAnthropic, decision.Platform)
require.Equal(t, "7d", decision.Window)
require.Empty(t, decision.Scope)
require.Equal(t, 82.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, wantUntil.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_OpenAIPreservesPercentageSemantics(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
openAIUntil := now.Add(24 * time.Hour)
openAIAccount := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_5h_used_percent": 1.0,
"codex_5h_reset_at": openAIUntil.Format(time.RFC3339),
},
}
candidate := openAIThresholdCandidate(openAIAccount.Extra, "5h", now)
require.NotNil(t, candidate)
require.Equal(t, 1.0, candidate.usedPercent)
openAIDecision := EvaluateAccountSchedulingThreshold(openAIAccount, map[string]int{
PlatformOpenAI: 90,
}, now)
require.False(t, openAIDecision.ShouldPause)
openAIAccount.Extra["codex_5h_used_percent"] = 91.0
openAIDecision = EvaluateAccountSchedulingThreshold(openAIAccount, map[string]int{
PlatformOpenAI: 90,
}, now)
require.True(t, openAIDecision.ShouldPause)
require.Equal(t, 91.0, openAIDecision.UsedPercent)
}
func TestEvaluateAccountSchedulingThreshold_OpenAISkipsStaleSnapshot(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_usage_updated_at": now.Add(-2 * time.Hour).Format(time.RFC3339),
"codex_5h_used_percent": 100.0,
"codex_5h_reset_at": now.Add(3 * time.Hour).Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
require.False(t, decision.ShouldPause)
}
func TestEvaluateAccountSchedulingThreshold_OpenAISkipsResetWindow(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
"codex_5h_used_percent": 100.0,
"codex_5h_reset_at": now.Add(-time.Second).Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
require.False(t, decision.ShouldPause)
}
func TestEvaluateAccountSchedulingThreshold_OpenAIPausesFreshExhaustedSnapshot(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
resetAt := now.Add(3 * time.Hour)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
"codex_5h_used_percent": 100.0,
"codex_5h_reset_at": resetAt.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, "5h", decision.Window)
require.Equal(t, 100.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, resetAt.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_OpenAIPausesFreshExhaustedSevenDayWindow(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
resetAt := now.Add(5 * 24 * time.Hour)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
"codex_7d_used_percent": 95.0,
"codex_7d_reset_at": resetAt.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, "7d", decision.Window)
require.Equal(t, 95.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, resetAt.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_AnthropicPreservesFractionalUtilizationSemantics(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
anthropicUntil := now.Add(5 * time.Hour)
anthropicAccount := &Account{
Platform: PlatformAnthropic,
SessionWindowEnd: &anthropicUntil,
Extra: map[string]any{
"session_window_utilization": 0.92,
},
}
anthropicDecision := EvaluateAccountSchedulingThreshold(anthropicAccount, map[string]int{
PlatformAnthropic: 90,
}, now)
require.True(t, anthropicDecision.ShouldPause)
require.Equal(t, 92.0, anthropicDecision.UsedPercent)
}
func TestEvaluateAccountSchedulingThreshold_AccountOverrideCanLowerOpenAIThreshold(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
wantUntil := now.Add(12 * time.Hour)
account := &Account{
Platform: PlatformOpenAI,
Credentials: map[string]any{
"account_scheduling_threshold": 80,
},
Extra: map[string]any{
"codex_7d_used_percent": 85.0,
"codex_7d_reset_at": wantUntil.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformOpenAI: 90,
}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, PlatformOpenAI, decision.Platform)
require.Equal(t, 80, decision.ThresholdPercent)
require.Equal(t, "7d", decision.Window)
require.Empty(t, decision.Scope)
require.Equal(t, 85.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, wantUntil.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_AccountOverrideHundredDisablesOpenAI(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
account := &Account{
Platform: PlatformOpenAI,
Credentials: map[string]any{
"account_scheduling_threshold": 100,
},
Extra: map[string]any{
"codex_7d_used_percent": 99.0,
"codex_7d_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformOpenAI: 80,
}, now)
require.False(t, decision.ShouldPause)
require.Equal(t, 100, decision.ThresholdPercent)
}
func TestEvaluateAccountSchedulingThreshold_AccountOverrideRoundsDecimalThreshold(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
wantUntil := now.Add(12 * time.Hour)
account := &Account{
Platform: PlatformOpenAI,
Credentials: map[string]any{
"account_scheduling_threshold": 75.5,
},
Extra: map[string]any{
"codex_7d_used_percent": 80.0,
"codex_7d_reset_at": wantUntil.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformOpenAI: 90,
}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, 76, decision.ThresholdPercent)
require.Equal(t, 80.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, wantUntil.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_UnsupportedPlatformsDoNotPause(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
cases := []struct {
name string
platform string
threshold int
extra map[string]any
}{
{
name: "gemini",
platform: PlatformGemini,
threshold: 80,
extra: map[string]any{
"gemini_usage_raw": map[string]any{
"buckets": []any{
map[string]any{
"modelId": "gemini-2.5-pro",
"remainingFraction": 0.05,
"resetTime": now.Add(2 * time.Hour).Format(time.RFC3339),
},
},
},
},
},
{
name: "kiro",
platform: PlatformKiro,
threshold: 90,
extra: map[string]any{
"kiro_sched_utilization": 99.0,
"kiro_sched_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339),
},
},
{
name: "antigravity",
platform: PlatformAntigravity,
threshold: 90,
extra: map[string]any{
"antigravity_sched_utilization": 92.0,
"antigravity_sched_reset_at": now.Add(48 * time.Hour).Format(time.RFC3339),
"antigravity_sched_scope": "gemini",
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
account := &Account{
Platform: tc.platform,
Credentials: map[string]any{
"account_scheduling_threshold": 1,
},
Extra: tc.extra,
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
tc.platform: tc.threshold,
}, now)
require.False(t, decision.ShouldPause)
require.Zero(t, decision.ThresholdPercent)
})
}
}
func TestEvaluateAccountSchedulingThreshold_GrokUsesConfiguredThresholds(t *testing.T) {
t.Parallel()
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
wantUntil := now.Add(2 * time.Hour)
account := &Account{
Platform: PlatformGrok,
Extra: map[string]any{
"grok_sched_utilization": 92.0,
"grok_sched_reset_at": wantUntil.Format(time.RFC3339),
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{
PlatformGrok: 90,
}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, PlatformGrok, decision.Platform)
require.Equal(t, 90, decision.ThresholdPercent)
require.Equal(t, "grok", decision.Scope)
require.Equal(t, 92.0, decision.UsedPercent)
require.NotNil(t, decision.Until)
require.True(t, wantUntil.Equal(*decision.Until))
}
func TestEvaluateAccountSchedulingThreshold_GrokUsesOnlyHeaderQuotaWindow(t *testing.T) {
t.Parallel()
// Billing seven_day/thirty_day must not drive pause; only grok_sched_* may.
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
weeklyEnd := now.Add(3 * time.Hour)
weeklyPct := 99.0
headerUntil := now.Add(2 * time.Hour)
account := &Account{
Platform: PlatformGrok,
Extra: map[string]any{
"grok_sched_utilization": 50.0, // below threshold
"grok_sched_reset_at": headerUntil.Format(time.RFC3339),
grokBillingExtraKey: &xai.BillingSummary{
UsagePercent: &weeklyPct,
PeriodEnd: weeklyEnd.Format(time.RFC3339),
},
},
}
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformGrok: 90}, now)
require.False(t, decision.ShouldPause, "high billing % alone must not pause under scheduling windows")
account.Extra["grok_sched_utilization"] = 95.0
decision = EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformGrok: 90}, now)
require.True(t, decision.ShouldPause)
require.Equal(t, "grok", decision.Scope)
require.Equal(t, "quota", decision.Window)
require.NotNil(t, decision.Until)
require.True(t, headerUntil.Equal(*decision.Until))
}
@@ -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)
}
@@ -0,0 +1,188 @@
package service
import (
"encoding/json"
"fmt"
"strings"
"time"
)
const AccountSchedulingThresholdReasonSource = "account_scheduling_threshold"
const (
defaultTempUnschedReasonErrorMessage = "temporary scheduling block reason unavailable"
defaultAccountSchedulingThresholdErrorMessage = "account scheduling threshold reached"
)
type tempUnschedReasonPayload struct {
Source string `json:"source,omitempty"`
Platform string `json:"platform,omitempty"`
Window string `json:"window,omitempty"`
Scope string `json:"scope,omitempty"`
ThresholdPercent int `json:"threshold_percent,omitempty"`
UsedPercent float64 `json:"used_percent,omitempty"`
UntilUnix int64 `json:"until_unix,omitempty"`
TriggeredAtUnix int64 `json:"triggered_at_unix,omitempty"`
ErrorMessage string `json:"error_message"`
}
type AccountSchedulingThresholdReasonInput struct {
Platform string
Window string
Scope string
ThresholdPercent int
UsedPercent float64
Until time.Time
Now time.Time
}
func BuildTempUnschedReasonPayload(source string, errorMessage string) string {
payload := tempUnschedReasonPayload{
Source: strings.TrimSpace(source),
ErrorMessage: normalizeTempUnschedReasonErrorMessage(errorMessage, defaultTempUnschedReasonErrorMessage),
}
raw, err := json.Marshal(payload)
if err != nil {
return payload.ErrorMessage
}
return string(raw)
}
func BuildAccountSchedulingThresholdReason(errorMessage string) string {
return BuildTempUnschedReasonPayload(
AccountSchedulingThresholdReasonSource,
normalizeTempUnschedReasonErrorMessage(errorMessage, defaultAccountSchedulingThresholdErrorMessage),
)
}
func BuildDetailedAccountSchedulingThresholdReason(input AccountSchedulingThresholdReasonInput) string {
triggeredAt := input.Now
if triggeredAt.IsZero() {
triggeredAt = time.Now().UTC()
}
payload := tempUnschedReasonPayload{
Source: AccountSchedulingThresholdReasonSource,
Platform: strings.TrimSpace(input.Platform),
Window: strings.TrimSpace(input.Window),
Scope: strings.TrimSpace(input.Scope),
ThresholdPercent: input.ThresholdPercent,
UsedPercent: input.UsedPercent,
TriggeredAtUnix: triggeredAt.Unix(),
ErrorMessage: buildAccountSchedulingThresholdErrorMessage(input),
}
if !input.Until.IsZero() {
payload.UntilUnix = input.Until.UTC().Unix()
}
raw, err := json.Marshal(payload)
if err != nil {
return payload.ErrorMessage
}
return string(raw)
}
func IsAccountSchedulingThresholdReason(rawReason string) bool {
payload, ok := parseTempUnschedReasonPayload(rawReason)
if !ok {
return false
}
return payload.Source == AccountSchedulingThresholdReasonSource
}
func parseTempUnschedReasonPayload(rawReason string) (tempUnschedReasonPayload, bool) {
rawReason = strings.TrimSpace(rawReason)
if rawReason == "" {
return tempUnschedReasonPayload{}, false
}
var payload tempUnschedReasonPayload
if err := json.Unmarshal([]byte(rawReason), &payload); err != nil {
return tempUnschedReasonPayload{}, false
}
payload.Source = strings.TrimSpace(payload.Source)
payload.ErrorMessage = strings.TrimSpace(payload.ErrorMessage)
return payload, true
}
func normalizeTempUnschedReasonErrorMessage(errorMessage string, fallback string) string {
errorMessage = strings.TrimSpace(errorMessage)
if errorMessage != "" {
return errorMessage
}
fallback = strings.TrimSpace(fallback)
if fallback != "" {
return fallback
}
return defaultTempUnschedReasonErrorMessage
}
func buildAccountSchedulingThresholdErrorMessage(input AccountSchedulingThresholdReasonInput) string {
platform := strings.TrimSpace(input.Platform)
if platform == "" {
platform = "account"
}
target := strings.TrimSpace(input.Window)
if scope := strings.TrimSpace(input.Scope); scope != "" {
if target == "" {
target = scope
} else {
target = target + "/" + scope
}
}
if target == "" {
target = "usage window"
}
threshold := input.ThresholdPercent
if threshold <= 0 {
threshold = 100
}
untilText := "the window reset"
if !input.Until.IsZero() {
untilText = input.Until.UTC().Format(time.RFC3339)
}
return fmt.Sprintf(
"%s scheduling threshold reached for %s: %.1f%% used >= %d%%; paused until %s",
platform,
target,
input.UsedPercent,
threshold,
untilText,
)
}
func tempUnschedStateFromStoredReason(rawReason string, fallbackUntilUnix int64) *TempUnschedState {
state := &TempUnschedState{
UntilUnix: fallbackUntilUnix,
RuleIndex: -1,
}
rawReason = strings.TrimSpace(rawReason)
if rawReason == "" {
state.ErrorMessage = defaultTempUnschedReasonErrorMessage
return state
}
parsed := TempUnschedState{RuleIndex: -1}
if err := json.Unmarshal([]byte(rawReason), &parsed); err == nil {
if fallbackUntilUnix > parsed.UntilUnix {
parsed.UntilUnix = fallbackUntilUnix
}
if strings.TrimSpace(parsed.ErrorMessage) == "" {
if IsAccountSchedulingThresholdReason(rawReason) {
parsed.ErrorMessage = defaultAccountSchedulingThresholdErrorMessage
} else {
parsed.ErrorMessage = rawReason
}
}
return &parsed
}
state.ErrorMessage = rawReason
return state
}
@@ -0,0 +1,76 @@
//go:build unit
package service
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestBuildAccountSchedulingThresholdReason_UsesSourceAndFallbackMessage(t *testing.T) {
raw := BuildAccountSchedulingThresholdReason(" \t ")
var payload map[string]string
require.NoError(t, json.Unmarshal([]byte(raw), &payload))
require.Equal(t, AccountSchedulingThresholdReasonSource, payload["source"])
require.Equal(t, defaultAccountSchedulingThresholdErrorMessage, payload["error_message"])
require.True(t, IsAccountSchedulingThresholdReason(raw))
}
func TestIsAccountSchedulingThresholdReason(t *testing.T) {
require.True(t, IsAccountSchedulingThresholdReason(BuildAccountSchedulingThresholdReason("threshold reached")))
require.False(t, IsAccountSchedulingThresholdReason(BuildTempUnschedReasonPayload("", "temporary block")))
require.False(t, IsAccountSchedulingThresholdReason("plain text reason"))
}
func TestTempUnschedStateFromStoredReason_EmptyReasonUsesFallbackErrorMessage(t *testing.T) {
state := tempUnschedStateFromStoredReason(" \n ", 1735689600)
require.NotNil(t, state)
require.Equal(t, int64(1735689600), state.UntilUnix)
require.Equal(t, defaultTempUnschedReasonErrorMessage, state.ErrorMessage)
}
func TestTempUnschedStateFromStoredReason_MissingRuleIndexIsSystemRule(t *testing.T) {
state := tempUnschedStateFromStoredReason(`{"error_message":"system cooldown"}`, 123)
require.Equal(t, -1, state.RuleIndex)
}
func TestTempUnschedStateFromStoredReason_SchedulingThresholdJSONWithoutMessageUsesThresholdFallback(t *testing.T) {
raw := `{"source":"` + AccountSchedulingThresholdReasonSource + `"}`
state := tempUnschedStateFromStoredReason(raw, 1735689600)
require.NotNil(t, state)
require.Equal(t, int64(1735689600), state.UntilUnix)
require.Equal(t, defaultAccountSchedulingThresholdErrorMessage, state.ErrorMessage)
}
func TestBuildDetailedAccountSchedulingThresholdReason_IncludesReadableFields(t *testing.T) {
now := time.Unix(1735689600, 0).UTC()
until := now.Add(5 * time.Hour)
raw := BuildDetailedAccountSchedulingThresholdReason(AccountSchedulingThresholdReasonInput{
Platform: PlatformOpenAI,
Window: "7d",
ThresholdPercent: 90,
UsedPercent: 92.5,
Until: until,
Now: now,
})
var payload map[string]any
require.NoError(t, json.Unmarshal([]byte(raw), &payload))
require.Equal(t, AccountSchedulingThresholdReasonSource, payload["source"])
require.Equal(t, PlatformOpenAI, payload["platform"])
require.Equal(t, "7d", payload["window"])
require.Equal(t, float64(90), payload["threshold_percent"])
require.Equal(t, float64(92.5), payload["used_percent"])
require.Equal(t, float64(until.Unix()), payload["until_unix"])
require.Equal(t, float64(now.Unix()), payload["triggered_at_unix"])
require.Contains(t, payload["error_message"], "openai scheduling threshold reached")
require.Contains(t, payload["error_message"], "92.5% used >= 90%")
}
+523
View File
@@ -0,0 +1,523 @@
package service
import (
"context"
"fmt"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
)
var (
ErrAccountNotFound = infraerrors.NotFound("ACCOUNT_NOT_FOUND", "account not found")
ErrAccountNilInput = infraerrors.BadRequest("ACCOUNT_NIL_INPUT", "account input cannot be nil")
ErrAccountNotInFallback = infraerrors.BadRequest("ACCOUNT_NOT_IN_FALLBACK", "account is not in proxy fallback state")
)
const AccountListGroupUngrouped int64 = -1
const AccountPrivacyModeUnsetFilter = "__unset__"
// OAuthRefreshPageOptions describes one bounded, cursor-stable scan of OAuth
// accounts. Candidate platforms are supplied by TokenRefreshService's refresher
// registry so repository eligibility cannot drift from registered providers.
type OAuthRefreshPageOptions struct {
Platforms []string
AfterID int64
Limit int
ActiveOnly bool
IncludeSetupToken bool
RequireRefreshToken bool
ExcludeRetryCooldown bool
}
// OAuthRefreshCandidatePage keeps cursor metadata from the raw SQL ID page.
// Hydration may legitimately lose a concurrently deleted row, but callers can
// still advance past the raw page without truncating or duplicating the scan.
type OAuthRefreshCandidatePage struct {
Accounts []Account
NextAfterID int64
HasMore bool
}
// OAuthRefreshCandidatePager is intentionally narrower than AccountRepository.
// Production refresh cycles fail closed when the repository does not implement
// this bounded contract instead of silently falling back to an unpaged scan.
type OAuthRefreshCandidatePager interface {
ListOAuthRefreshCandidatePage(ctx context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error)
}
type AccountRepository interface {
Create(ctx context.Context, account *Account) error
GetByID(ctx context.Context, id int64) (*Account, error)
// GetByIDs fetches accounts by IDs in a single query.
// It should return all accounts found (missing IDs are ignored).
GetByIDs(ctx context.Context, ids []int64) ([]*Account, error)
// ExistsByID 检查账号是否存在,仅返回布尔值,用于删除前的轻量级存在性检查
ExistsByID(ctx context.Context, id int64) (bool, error)
// GetByCRSAccountID finds an account previously synced from CRS.
// Returns (nil, nil) if not found.
GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error)
// FindByExtraField 根据 extra 字段中的键值对查找账号
FindByExtraField(ctx context.Context, key string, value any) ([]Account, error)
// ListCRSAccountIDs returns a map of crs_account_id -> local account ID
// for all accounts that have been synced from CRS.
ListCRSAccountIDs(ctx context.Context) (map[string]int64, error)
Update(ctx context.Context, account *Account) error
Delete(ctx context.Context, id int64) error
List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error)
ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error)
// ListAllWithFilters 返回符合过滤条件的全部账号(不分页),用于账号列表页
// 计算 OpenAI 调度分数的过滤范围池。
ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error)
ListByGroup(ctx context.Context, groupID int64) ([]Account, error)
ListActive(ctx context.Context) ([]Account, error)
ListByPlatform(ctx context.Context, platform string) ([]Account, error)
UpdateLastUsed(ctx context.Context, id int64) error
BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error
SetError(ctx context.Context, id int64, errorMsg string) error
ClearError(ctx context.Context, id int64) error
SetSchedulable(ctx context.Context, id int64, schedulable bool) error
AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error)
BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error
ListSchedulable(ctx context.Context) ([]Account, error)
ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error)
ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error)
ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error)
ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error)
ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error)
ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error)
// ListModelAvailabilityCandidates returns accounts that are enabled by
// persistent configuration (active + schedulable) for model-support
// diagnosis. It deliberately does not filter transient runtime state such
// as rate-limit, overload, temporary-unschedulable, or expiry windows.
// When groupID is nil, includeGrouped controls whether the query scans all
// matching accounts or only accounts without a group binding.
ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error)
SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error
SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error
SetOverloaded(ctx context.Context, id int64, until time.Time) error
SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error
ClearTempUnschedulable(ctx context.Context, id int64) error
ClearRateLimit(ctx context.Context, id int64) error
ClearAntigravityQuotaScopes(ctx context.Context, id int64) error
ClearModelRateLimits(ctx context.Context, id int64) error
UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error
// UpdateSessionWindowEnd 仅更新 5h 窗口的结束时间,不动 start / status。
// 用于 active poll 拿到新 ResetsAt 后回写,避免覆盖请求路径上记录的 status。
UpdateSessionWindowEnd(ctx context.Context, id int64, end time.Time) error
UpdateExtra(ctx context.Context, id int64, updates map[string]any) error
BulkUpdate(ctx context.Context, ids []int64, updates AccountBulkUpdate) (int64, error)
// IncrementQuotaUsed 原子递增 API Key 账号的配额用量(总/日/周)
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) error
// ResetQuotaUsed 重置 API Key 账号所有维度的配额用量为 0
ResetQuotaUsed(ctx context.Context, id int64) error
// RevertProxyFallback 将账号的 proxy_id 切回 proxy_fallback_origin_id,并清空 origin 字段。
// 仅当 proxy_fallback_origin_id IS NOT NULL 时更新,否则视为账号不存在(返回 ErrAccountNotFound)。
RevertProxyFallback(ctx context.Context, accountID int64) error
// ListShadowsByParent 返回指定父账号的影子账号;当前实现仅查 quota_dimension='spark'(唯一预设)。
// ⚠️ 新增影子维度时:须更新此函数(或新增维度专用列举),并检查所有调用点(级联删除/一母一影校验/type 守卫),否则会静默漏掉新维度。
ListShadowsByParent(ctx context.Context, parentID int64) ([]*Account, error)
}
type AccountDuplicateRepository interface {
// CreateWithAccountGroups atomically persists an account, its exact group priorities,
// and the scheduler outbox event for the new routing snapshot.
CreateWithAccountGroups(ctx context.Context, account *Account, groups []AccountGroup) error
}
// AccountBillingSettingsRepository applies an admin edit without overwriting a
// rate_multiplier that a successful upstream probe synchronized after the edit
// form was loaded. A nil rateMultiplier means the request did not edit it.
type AccountBillingSettingsRepository interface {
UpdateWithAccountBillingSettings(
ctx context.Context,
account *Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error
}
// AdminAccountRepository makes the account-duplication write capability an explicit
// construction dependency without forcing read-only gateway test doubles to implement it.
type AdminAccountRepository interface {
AccountRepository
AccountDuplicateRepository
AccountBillingSettingsRepository
}
// AccountBulkUpdate describes the fields that can be updated in a bulk operation.
// Nil pointers mean "do not change".
type AccountBulkUpdate struct {
Name *string
ProxyID *int64
Concurrency *int
Priority *int
RateMultiplier *float64
LoadFactor *int
Status *string
Schedulable *bool
Credentials map[string]any
Extra map[string]any
ProbeEnabled *bool
// EnsureCodexFingerprintSeed asks the repository to atomically preserve an
// existing valid Codex fingerprint seed or create one for eligible rows.
EnsureCodexFingerprintSeed bool
}
// CreateAccountRequest 创建账号请求
type CreateAccountRequest struct {
Name string `json:"name"`
Notes *string `json:"notes"`
Platform string `json:"platform"`
Type string `json:"type"`
Credentials map[string]any `json:"credentials"`
Extra map[string]any `json:"extra"`
ProxyID *int64 `json:"proxy_id"`
Concurrency int `json:"concurrency"`
Priority int `json:"priority"`
GroupIDs []int64 `json:"group_ids"`
ExpiresAt *time.Time `json:"expires_at"`
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
}
// UpdateAccountRequest 更新账号请求
type UpdateAccountRequest struct {
Name *string `json:"name"`
Notes *string `json:"notes"`
Credentials *map[string]any `json:"credentials"`
Extra *map[string]any `json:"extra"`
ProxyID *int64 `json:"proxy_id"`
Concurrency *int `json:"concurrency"`
Priority *int `json:"priority"`
Status *string `json:"status"`
GroupIDs *[]int64 `json:"group_ids"`
ExpiresAt *time.Time `json:"expires_at"`
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
}
// AccountService 账号管理服务
type AccountService struct {
accountRepo AccountRepository
groupRepo GroupRepository
}
type groupExistenceBatchChecker interface {
ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error)
}
// NewAccountService 创建账号服务实例
func NewAccountService(accountRepo AccountRepository, groupRepo GroupRepository) *AccountService {
return &AccountService{
accountRepo: accountRepo,
groupRepo: groupRepo,
}
}
// Create 创建账号
func (s *AccountService) Create(ctx context.Context, req CreateAccountRequest) (*Account, error) {
// 验证分组是否存在(如果指定了分组)
if len(req.GroupIDs) > 0 {
if err := s.validateGroupIDsExist(ctx, req.GroupIDs); err != nil {
return nil, err
}
}
// 创建账号
account := &Account{
Name: req.Name,
Notes: normalizeAccountNotes(req.Notes),
Platform: req.Platform,
Type: req.Type,
Credentials: SanitizeStoredCredentials(req.Platform, req.Credentials),
Extra: prepareCodexFingerprintExtraForCreate(req.Platform, req.Type, req.Extra),
ProxyID: req.ProxyID,
Concurrency: req.Concurrency,
Priority: req.Priority,
Status: StatusActive,
ExpiresAt: req.ExpiresAt,
}
if req.AutoPauseOnExpired != nil {
account.AutoPauseOnExpired = *req.AutoPauseOnExpired
} else {
account.AutoPauseOnExpired = true
}
if err := s.accountRepo.Create(ctx, account); err != nil {
return nil, fmt.Errorf("create account: %w", err)
}
// require_oauth_only 检查:apikey 类型账号不可加入限制分组
if account.Type == AccountTypeAPIKey && len(req.GroupIDs) > 0 {
for _, gid := range req.GroupIDs {
g, err := s.groupRepo.GetByID(ctx, gid)
if err != nil {
return nil, err
}
if g.RequireOAuthOnly && (g.Platform == PlatformOpenAI || g.Platform == PlatformAntigravity || g.Platform == PlatformAnthropic || g.Platform == PlatformGemini || g.Platform == PlatformGrok) {
return nil, fmt.Errorf("分组 [%s] 仅允许 OAuth 账号,apikey 类型账号无法加入", g.Name)
}
}
}
// 绑定分组
if len(req.GroupIDs) > 0 {
if err := s.accountRepo.BindGroups(ctx, account.ID, req.GroupIDs); err != nil {
return nil, fmt.Errorf("bind groups: %w", err)
}
}
return account, nil
}
// GetByID 根据ID获取账号
func (s *AccountService) GetByID(ctx context.Context, id int64) (*Account, error) {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
return nil, fmt.Errorf("get account: %w", err)
}
return account, nil
}
// List 获取账号列表
func (s *AccountService) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
accounts, pagination, err := s.accountRepo.List(ctx, params)
if err != nil {
return nil, nil, fmt.Errorf("list accounts: %w", err)
}
return accounts, pagination, nil
}
// ListByPlatform 根据平台获取账号列表
func (s *AccountService) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
accounts, err := s.accountRepo.ListByPlatform(ctx, platform)
if err != nil {
return nil, fmt.Errorf("list accounts by platform: %w", err)
}
return accounts, nil
}
// ListByGroup 根据分组获取账号列表
func (s *AccountService) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
accounts, err := s.accountRepo.ListByGroup(ctx, groupID)
if err != nil {
return nil, fmt.Errorf("list accounts by group: %w", err)
}
return accounts, nil
}
// Update 更新账号
func (s *AccountService) Update(ctx context.Context, id int64, req UpdateAccountRequest) (*Account, error) {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
return nil, fmt.Errorf("get account: %w", err)
}
// 更新字段
if req.Name != nil {
account.Name = *req.Name
}
if req.Notes != nil {
account.Notes = normalizeAccountNotes(req.Notes)
}
if req.Credentials != nil {
account.Credentials = SanitizeStoredCredentials(account.Platform, *req.Credentials)
}
if req.Extra != nil {
extra := make(map[string]any, len(*req.Extra))
for key, value := range *req.Extra {
extra[key] = value
}
delete(extra, OllamaCloudUsageSessionExtraKey)
delete(extra, OllamaCloudUsageAutoRefreshExtraKey)
delete(extra, OllamaCloudUsageSnapshotExtraKey)
account.Extra = prepareCodexFingerprintExtraForUpdate(account, extra)
} else {
account.Extra = prepareCodexFingerprintExtraForUpdate(account, account.Extra)
}
if req.ProxyID != nil {
account.ProxyID = req.ProxyID
}
if req.Concurrency != nil {
account.Concurrency = *req.Concurrency
}
if req.Priority != nil {
account.Priority = *req.Priority
}
if req.Status != nil {
account.Status = *req.Status
}
if req.ExpiresAt != nil {
account.ExpiresAt = req.ExpiresAt
}
if req.AutoPauseOnExpired != nil {
account.AutoPauseOnExpired = *req.AutoPauseOnExpired
}
// 先验证分组是否存在(在任何写操作之前)
if req.GroupIDs != nil {
if err := s.validateGroupIDsExist(ctx, *req.GroupIDs); err != nil {
return nil, err
}
}
// 执行更新
if err := s.accountRepo.Update(ctx, account); err != nil {
return nil, fmt.Errorf("update account: %w", err)
}
// require_oauth_only 检查
if account.Type == AccountTypeAPIKey && req.GroupIDs != nil {
for _, gid := range *req.GroupIDs {
g, err := s.groupRepo.GetByID(ctx, gid)
if err != nil {
return nil, err
}
if g.RequireOAuthOnly && (g.Platform == PlatformOpenAI || g.Platform == PlatformAntigravity || g.Platform == PlatformAnthropic || g.Platform == PlatformGemini || g.Platform == PlatformGrok) {
return nil, fmt.Errorf("分组 [%s] 仅允许 OAuth 账号,apikey 类型账号无法加入", g.Name)
}
}
}
// 绑定分组
if req.GroupIDs != nil {
if err := s.accountRepo.BindGroups(ctx, account.ID, *req.GroupIDs); err != nil {
return nil, fmt.Errorf("bind groups: %w", err)
}
}
return account, nil
}
// Delete 删除账号
// 优化:使用 ExistsByID 替代 GetByID 进行存在性检查,
// 避免加载完整账号对象及其关联数据,提升删除操作的性能
func (s *AccountService) Delete(ctx context.Context, id int64) error {
// 使用轻量级的存在性检查,而非加载完整账号对象
exists, err := s.accountRepo.ExistsByID(ctx, id)
if err != nil {
return fmt.Errorf("check account: %w", err)
}
// 明确返回账号不存在错误,便于调用方区分错误类型
if !exists {
return ErrAccountNotFound
}
// 注意:此处不级联删除 spark 影子账号。当前唯一的后台删除入口走 AdminService.DeleteAccount
// (已 ListShadowsByParent 先删影子再删母)。本方法目前无删除调用方;若未来有调用方经此
// 删除母账号,需在此补级联,否则会留下孤儿影子(外审第6轮 P3:当前不可达,记为残留)。
if err := s.accountRepo.Delete(ctx, id); err != nil {
return fmt.Errorf("delete account: %w", err)
}
return nil
}
func (s *AccountService) validateGroupIDsExist(ctx context.Context, groupIDs []int64) error {
if len(groupIDs) == 0 {
return nil
}
if s.groupRepo == nil {
return fmt.Errorf("group repository not configured")
}
if batchChecker, ok := s.groupRepo.(groupExistenceBatchChecker); ok {
existsByID, err := batchChecker.ExistsByIDs(ctx, groupIDs)
if err != nil {
return fmt.Errorf("check groups exists: %w", err)
}
for _, groupID := range groupIDs {
if groupID <= 0 {
return fmt.Errorf("get group: %w", ErrGroupNotFound)
}
if !existsByID[groupID] {
return fmt.Errorf("get group: %w", ErrGroupNotFound)
}
}
return nil
}
for _, groupID := range groupIDs {
_, err := s.groupRepo.GetByID(ctx, groupID)
if err != nil {
return fmt.Errorf("get group: %w", err)
}
}
return nil
}
// UpdateStatus 更新账号状态
func (s *AccountService) UpdateStatus(ctx context.Context, id int64, status string, errorMessage string) error {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
return fmt.Errorf("get account: %w", err)
}
account.Status = status
account.ErrorMessage = errorMessage
if err := s.accountRepo.Update(ctx, account); err != nil {
return fmt.Errorf("update account: %w", err)
}
return nil
}
// UpdateLastUsed 更新最后使用时间
func (s *AccountService) UpdateLastUsed(ctx context.Context, id int64) error {
if err := s.accountRepo.UpdateLastUsed(ctx, id); err != nil {
return fmt.Errorf("update last used: %w", err)
}
return nil
}
// GetCredential 获取账号凭证(安全访问)
func (s *AccountService) GetCredential(ctx context.Context, id int64, key string) (string, error) {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
return "", fmt.Errorf("get account: %w", err)
}
return account.GetCredential(key), nil
}
// TestCredentials 测试账号凭证是否有效(需要实现具体平台的测试逻辑)
func (s *AccountService) TestCredentials(ctx context.Context, id int64) error {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
return fmt.Errorf("get account: %w", err)
}
// 根据平台执行不同的测试逻辑
switch account.Platform {
case PlatformAnthropic:
// TODO: 测试Anthropic API凭证
return nil
case PlatformOpenAI:
// TODO: 测试OpenAI API凭证
return nil
case PlatformGemini:
// TODO: 测试Gemini API凭证
return nil
case PlatformGrok:
// Grok OAuth credentials are validated via token exchange/refresh and request-path probes.
return nil
case PlatformKimi, PlatformZhipu, PlatformDeepseek:
// 国产 OpenAI 兼容供应商:凭证为 API Key,实际可用性经余额/额度探测与转发路径验证。
return nil
default:
return fmt.Errorf("unsupported platform: %s", account.Platform)
}
}
@@ -0,0 +1,291 @@
//go:build unit
// 账号服务删除方法的单元测试
// 测试 AccountService.Delete 方法在各种场景下的行为
package service
import (
"context"
"errors"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
// accountRepoStub 是 AccountRepository 接口的测试桩实现。
// 用于隔离测试 AccountService.Delete 方法,避免依赖真实数据库。
//
// 设计说明:
// - exists: 模拟 ExistsByID 返回的存在性结果
// - existsErr: 模拟 ExistsByID 返回的错误
// - deleteErr: 模拟 Delete 返回的错误
// - deletedIDs: 记录被调用删除的账号 ID,用于断言验证
type accountRepoStub struct {
exists bool // ExistsByID 的返回值
existsErr error // ExistsByID 的错误返回值
deleteErr error // Delete 的错误返回值
deletedIDs []int64 // 记录已删除的账号 ID 列表
}
// 以下方法在本测试中不应被调用,使用 panic 确保测试失败时能快速定位问题
func (s *accountRepoStub) Create(ctx context.Context, account *Account) error {
panic("unexpected Create call")
}
func (s *accountRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
panic("unexpected GetByID call")
}
func (s *accountRepoStub) GetByIDs(ctx context.Context, ids []int64) ([]*Account, error) {
panic("unexpected GetByIDs call")
}
// ExistsByID 返回预设的存在性检查结果。
// 这是 Delete 方法调用的第一个仓储方法,用于验证账号是否存在。
func (s *accountRepoStub) ExistsByID(ctx context.Context, id int64) (bool, error) {
return s.exists, s.existsErr
}
func (s *accountRepoStub) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
panic("unexpected GetByCRSAccountID call")
}
func (s *accountRepoStub) FindByExtraField(ctx context.Context, key string, value any) ([]Account, error) {
panic("unexpected FindByExtraField call")
}
func (s *accountRepoStub) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
panic("unexpected ListCRSAccountIDs call")
}
func (s *accountRepoStub) Update(ctx context.Context, account *Account) error {
panic("unexpected Update call")
}
// Delete 记录被删除的账号 ID 并返回预设的错误。
// 通过 deletedIDs 可以验证删除操作是否被正确调用。
func (s *accountRepoStub) Delete(ctx context.Context, id int64) error {
s.deletedIDs = append(s.deletedIDs, id)
return s.deleteErr
}
// 以下是接口要求实现但本测试不关心的方法
func (s *accountRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *accountRepoStub) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) {
return nil, nil
}
func (s *accountRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *accountRepoStub) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
panic("unexpected ListByGroup call")
}
func (s *accountRepoStub) ListActive(ctx context.Context) ([]Account, error) {
panic("unexpected ListActive call")
}
func (s *accountRepoStub) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
panic("unexpected ListByPlatform call")
}
func (s *accountRepoStub) UpdateLastUsed(ctx context.Context, id int64) error {
panic("unexpected UpdateLastUsed call")
}
func (s *accountRepoStub) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
panic("unexpected BatchUpdateLastUsed call")
}
func (s *accountRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
panic("unexpected SetError call")
}
func (s *accountRepoStub) ClearError(ctx context.Context, id int64) error {
panic("unexpected ClearError call")
}
func (s *accountRepoStub) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
panic("unexpected SetSchedulable call")
}
func (s *accountRepoStub) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
panic("unexpected AutoPauseExpiredAccounts call")
}
func (s *accountRepoStub) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
panic("unexpected BindGroups call")
}
func (s *accountRepoStub) ListSchedulable(ctx context.Context) ([]Account, error) {
panic("unexpected ListSchedulable call")
}
func (s *accountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
panic("unexpected ListSchedulableByGroupID call")
}
func (s *accountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
panic("unexpected ListSchedulableByPlatform call")
}
func (s *accountRepoStub) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
panic("unexpected ListSchedulableByGroupIDAndPlatform call")
}
func (s *accountRepoStub) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
panic("unexpected ListSchedulableByPlatforms call")
}
func (s *accountRepoStub) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) {
panic("unexpected ListSchedulableByGroupIDAndPlatforms call")
}
func (s *accountRepoStub) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
panic("unexpected ListSchedulableUngroupedByPlatform call")
}
func (s *accountRepoStub) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
panic("unexpected ListSchedulableUngroupedByPlatforms call")
}
func (s *accountRepoStub) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) {
panic("unexpected ListModelAvailabilityCandidates call")
}
func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
panic("unexpected SetRateLimited call")
}
func (s *accountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
panic("unexpected SetModelRateLimit call")
}
func (s *accountRepoStub) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
panic("unexpected SetOverloaded call")
}
func (s *accountRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
panic("unexpected SetTempUnschedulable call")
}
func (s *accountRepoStub) ClearTempUnschedulable(ctx context.Context, id int64) error {
panic("unexpected ClearTempUnschedulable call")
}
func (s *accountRepoStub) ClearRateLimit(ctx context.Context, id int64) error {
panic("unexpected ClearRateLimit call")
}
func (s *accountRepoStub) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
panic("unexpected ClearAntigravityQuotaScopes call")
}
func (s *accountRepoStub) ClearModelRateLimits(ctx context.Context, id int64) error {
panic("unexpected ClearModelRateLimits call")
}
func (s *accountRepoStub) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
panic("unexpected UpdateSessionWindow call")
}
func (s *accountRepoStub) UpdateSessionWindowEnd(ctx context.Context, id int64, end time.Time) error {
panic("unexpected UpdateSessionWindowEnd call")
}
func (s *accountRepoStub) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
panic("unexpected UpdateExtra call")
}
func (s *accountRepoStub) BulkUpdate(ctx context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
panic("unexpected BulkUpdate call")
}
func (s *accountRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) error {
return nil
}
func (s *accountRepoStub) ResetQuotaUsed(ctx context.Context, id int64) error {
return nil
}
func (s *accountRepoStub) RevertProxyFallback(ctx context.Context, accountID int64) error {
panic("unexpected RevertProxyFallback call")
}
func (s *accountRepoStub) ListShadowsByParent(ctx context.Context, parentID int64) ([]*Account, error) {
return nil, nil
}
// TestAccountService_Delete_NotFound 测试删除不存在的账号时返回正确的错误。
// 预期行为:
// - ExistsByID 返回 false(账号不存在)
// - 返回 ErrAccountNotFound 错误
// - Delete 方法不被调用(deletedIDs 为空)
func TestAccountService_Delete_NotFound(t *testing.T) {
repo := &accountRepoStub{exists: false}
svc := &AccountService{accountRepo: repo}
err := svc.Delete(context.Background(), 55)
require.ErrorIs(t, err, ErrAccountNotFound)
require.Empty(t, repo.deletedIDs) // 验证删除操作未被调用
}
// TestAccountService_Delete_CheckError 测试存在性检查失败时的错误处理。
// 预期行为:
// - ExistsByID 返回数据库错误
// - 返回包含 "check account" 的错误信息
// - Delete 方法不被调用
func TestAccountService_Delete_CheckError(t *testing.T) {
repo := &accountRepoStub{existsErr: errors.New("db down")}
svc := &AccountService{accountRepo: repo}
err := svc.Delete(context.Background(), 55)
require.Error(t, err)
require.ErrorContains(t, err, "check account") // 验证错误信息包含上下文
require.Empty(t, repo.deletedIDs)
}
// TestAccountService_Delete_DeleteError 测试删除操作失败时的错误处理。
// 预期行为:
// - ExistsByID 返回 true(账号存在)
// - Delete 被调用但返回错误
// - 返回包含 "delete account" 的错误信息
// - deletedIDs 记录了尝试删除的 ID
func TestAccountService_Delete_DeleteError(t *testing.T) {
repo := &accountRepoStub{
exists: true,
deleteErr: errors.New("delete failed"),
}
svc := &AccountService{accountRepo: repo}
err := svc.Delete(context.Background(), 55)
require.Error(t, err)
require.ErrorContains(t, err, "delete account")
require.Equal(t, []int64{55}, repo.deletedIDs) // 验证删除操作被调用
}
// TestAccountService_Delete_Success 测试删除操作成功的场景。
// 预期行为:
// - ExistsByID 返回 true(账号存在)
// - Delete 成功执行
// - 返回 nil 错误
// - deletedIDs 记录了被删除的 ID
func TestAccountService_Delete_Success(t *testing.T) {
repo := &accountRepoStub{exists: true}
svc := &AccountService{accountRepo: repo}
err := svc.Delete(context.Background(), 55)
require.NoError(t, err)
require.Equal(t, []int64{55}, repo.deletedIDs) // 验证正确的 ID 被删除
}
@@ -0,0 +1,19 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAccountSparkShadowHelpers(t *testing.T) {
pid := int64(100)
normal := &Account{ID: 100}
require.False(t, normal.IsShadow())
require.False(t, normal.IsCredentialShadow())
require.Equal(t, QuotaDimensionGlobal, normal.QuotaDimensionOrDefault())
shadow := &Account{ID: 200, ParentAccountID: &pid, QuotaDimension: QuotaDimensionSpark}
require.True(t, shadow.IsShadow())
require.True(t, shadow.IsCredentialShadow())
require.Equal(t, QuotaDimensionSpark, shadow.QuotaDimensionOrDefault())
}
@@ -0,0 +1,255 @@
package service
import (
"context"
"strings"
)
// resolveAccountStatsCost 计算账号统计定价费用。
// 返回 nil 表示不覆盖,使用默认公式(total_cost × account_rate_multiplier)。
//
// 优先级(先命中为准):
// 1. 自定义规则(始终尝试,不依赖 ApplyPricingToAccountStats 开关)
// 2. ApplyPricingToAccountStats 启用时,直接使用本次请求的客户计费(倍率前的 totalCost)
// 3. 模型定价文件(LiteLLM)中上游模型的默认价格
// 4. nil → 走默认公式(total_cost × account_rate_multiplier
//
// upstreamModel 是最终发往上游的模型 ID。
// totalCost 是本次请求的客户计费(倍率前),用于优先级 2。
// serviceTier 是最终参与用户计费的 OpenAI 服务层级,用于优先级 3。
func resolveAccountStatsCost(
ctx context.Context,
channelService *ChannelService,
billingService *BillingService,
accountID int64,
groupID int64,
upstreamModel string,
tokens UsageTokens,
requestCount int,
totalCost float64,
serviceTier string,
) *float64 {
if channelService == nil || upstreamModel == "" {
return nil
}
channel, err := channelService.GetChannelForGroup(ctx, groupID)
if err != nil || channel == nil {
return nil
}
platform := channelService.GetGroupPlatform(ctx, groupID)
// 优先级 1:自定义规则(始终尝试)
if cost := tryCustomRules(channel, accountID, groupID, platform, upstreamModel, tokens, requestCount); cost != nil {
return cost
}
// 优先级 2:渠道开启"应用模型定价到账号统计"时,直接使用客户计费(倍率前)
if channel.ApplyPricingToAccountStats {
cost := totalCost
if cost <= 0 {
return nil
}
return &cost
}
// 优先级 3:模型定价文件(LiteLLM)默认价格
if billingService != nil {
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier)
}
return nil
}
// tryModelFilePricing 使用模型定价文件(LiteLLM/fallback)中的价格计算费用。
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string) *float64 {
pricing, err := billingService.GetModelPricing(model)
if err != nil || pricing == nil {
return nil
}
normalizedTier := normalizeBillingServiceTier(serviceTier)
if normalizedTier == "priority" || normalizedTier == "fast" || normalizedTier == "flex" ||
billingService.shouldApplySessionLongContextPricing(tokens, pricing) {
breakdown, err := billingService.CalculateCostWithServiceTier(model, tokens, 1, normalizedTier)
if err != nil || breakdown == nil || breakdown.TotalCost <= 0 {
return nil
}
return &breakdown.TotalCost
}
cost := float64(tokens.InputTokens)*pricing.InputPricePerToken +
float64(tokens.OutputTokens)*pricing.OutputPricePerToken +
float64(tokens.CacheCreationTokens)*pricing.CacheCreationPricePerToken +
float64(tokens.CacheReadTokens)*pricing.CacheReadPricePerToken +
float64(tokens.ImageOutputTokens)*pricing.ImageOutputPricePerToken
if cost <= 0 {
return nil
}
return &cost
}
// tryCustomRules 遍历自定义规则,按数组顺序先命中为准。
func tryCustomRules(
channel *Channel, accountID, groupID int64,
platform, model string, tokens UsageTokens, requestCount int,
) *float64 {
modelLower := strings.ToLower(model)
for _, rule := range channel.AccountStatsPricingRules {
if !matchAccountStatsRule(&rule, accountID, groupID) {
continue
}
pricing := findPricingForModel(rule.Pricing, platform, modelLower)
if pricing == nil {
continue // 规则匹配但模型不在规则定价中,继续下一条
}
return calculateStatsCost(pricing, tokens, requestCount)
}
return nil
}
// matchAccountStatsRule 检查规则是否匹配指定的 accountID 和 groupID。
// 匹配条件:accountID ∈ rule.AccountIDs 或 groupID ∈ rule.GroupIDs。
// 如果规则的 AccountIDs 和 GroupIDs 都为空,视为不匹配。
func matchAccountStatsRule(rule *AccountStatsPricingRule, accountID, groupID int64) bool {
if len(rule.AccountIDs) == 0 && len(rule.GroupIDs) == 0 {
return false
}
for _, id := range rule.AccountIDs {
if id == accountID {
return true
}
}
for _, id := range rule.GroupIDs {
if id == groupID {
return true
}
}
return false
}
// findPricingForModel 在定价列表中查找匹配的模型定价。
// 先精确匹配,再通配符匹配(按配置顺序,先匹配先使用)。
func findPricingForModel(pricingList []ChannelModelPricing, platform, modelLower string) *ChannelModelPricing {
// 精确匹配优先
for i := range pricingList {
p := &pricingList[i]
if !isPlatformMatch(platform, p.Platform) {
continue
}
for _, m := range p.Models {
if strings.ToLower(m) == modelLower {
return p
}
}
}
// 通配符匹配:按配置顺序,先匹配先使用
for i := range pricingList {
p := &pricingList[i]
if !isPlatformMatch(platform, p.Platform) {
continue
}
for _, m := range p.Models {
ml := strings.ToLower(m)
if !strings.HasSuffix(ml, "*") {
continue
}
prefix := strings.TrimSuffix(ml, "*")
if strings.HasPrefix(modelLower, prefix) {
return p
}
}
}
return nil
}
// isPlatformMatch 判断平台是否匹配(空平台视为不限平台)。
func isPlatformMatch(queryPlatform, pricingPlatform string) bool {
if queryPlatform == "" || pricingPlatform == "" {
return true
}
return queryPlatform == pricingPlatform
}
// calculateStatsCost 使用给定的定价计算费用(不含任何倍率,原始费用)。
func calculateStatsCost(pricing *ChannelModelPricing, tokens UsageTokens, requestCount int) *float64 {
if pricing == nil {
return nil
}
switch pricing.BillingMode {
case BillingModePerRequest, BillingModeImage:
return calculatePerRequestStatsCost(pricing, requestCount)
default:
return calculateTokenStatsCost(pricing, tokens)
}
}
// calculatePerRequestStatsCost 按次/图片计费。
func calculatePerRequestStatsCost(pricing *ChannelModelPricing, requestCount int) *float64 {
if pricing.PerRequestPrice == nil || *pricing.PerRequestPrice <= 0 {
return nil
}
cost := *pricing.PerRequestPrice * float64(requestCount)
return &cost
}
// calculateTokenStatsCost Token 计费。
// If the pricing has intervals, find the matching interval by total token count
// and use its prices instead of the flat pricing fields.
func calculateTokenStatsCost(pricing *ChannelModelPricing, tokens UsageTokens) *float64 {
p := pricing
if len(pricing.Intervals) > 0 {
totalTokens := tokens.InputTokens + tokens.OutputTokens + tokens.CacheCreationTokens + tokens.CacheReadTokens
if iv := FindMatchingInterval(pricing.Intervals, totalTokens); iv != nil {
p = &ChannelModelPricing{
InputPrice: iv.InputPrice,
OutputPrice: iv.OutputPrice,
CacheWritePrice: iv.CacheWritePrice,
CacheReadPrice: iv.CacheReadPrice,
PerRequestPrice: iv.PerRequestPrice,
}
}
}
deref := func(ptr *float64) float64 {
if ptr == nil {
return 0
}
return *ptr
}
cost := float64(tokens.InputTokens)*deref(p.InputPrice) +
float64(tokens.OutputTokens)*deref(p.OutputPrice) +
float64(tokens.CacheCreationTokens)*deref(p.CacheWritePrice) +
float64(tokens.CacheReadTokens)*deref(p.CacheReadPrice) +
float64(tokens.ImageOutputTokens)*deref(p.ImageOutputPrice)
if cost <= 0 {
return nil
}
return &cost
}
// applyAccountStatsCost resolves the account stats cost for a usage log entry.
// It resolves the upstream model (falling back to the requested model) and calls
// the 4-level priority chain via resolveAccountStatsCost.
func applyAccountStatsCost(
ctx context.Context,
usageLog *UsageLog,
cs *ChannelService, bs *BillingService,
accountID int64, groupID int64,
upstreamModel, requestedModel string,
tokens UsageTokens,
totalCost float64,
) {
model := upstreamModel
if model == "" {
model = requestedModel
}
requestCount := 1
if usageLog != nil && usageLog.ImageCount > 0 {
requestCount = usageLog.ImageCount
}
serviceTier := ""
if usageLog != nil && usageLog.ServiceTier != nil {
serviceTier = *usageLog.ServiceTier
}
usageLog.AccountStatsCost = resolveAccountStatsCost(
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier,
)
}
@@ -0,0 +1,928 @@
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// matchAccountStatsRule
// ---------------------------------------------------------------------------
func TestMatchAccountStatsRule_BothEmpty_NoMatch(t *testing.T) {
rule := &AccountStatsPricingRule{}
require.False(t, matchAccountStatsRule(rule, 1, 10))
}
func TestMatchAccountStatsRule_AccountIDMatch(t *testing.T) {
rule := &AccountStatsPricingRule{AccountIDs: []int64{1, 2, 3}}
require.True(t, matchAccountStatsRule(rule, 2, 999))
}
func TestMatchAccountStatsRule_GroupIDMatch(t *testing.T) {
rule := &AccountStatsPricingRule{GroupIDs: []int64{10, 20}}
require.True(t, matchAccountStatsRule(rule, 999, 20))
}
func TestMatchAccountStatsRule_BothConfigured_AccountMatch(t *testing.T) {
rule := &AccountStatsPricingRule{
AccountIDs: []int64{1, 2},
GroupIDs: []int64{10, 20},
}
require.True(t, matchAccountStatsRule(rule, 2, 999))
}
func TestMatchAccountStatsRule_BothConfigured_GroupMatch(t *testing.T) {
rule := &AccountStatsPricingRule{
AccountIDs: []int64{1, 2},
GroupIDs: []int64{10, 20},
}
require.True(t, matchAccountStatsRule(rule, 999, 10))
}
func TestMatchAccountStatsRule_BothConfigured_NeitherMatch(t *testing.T) {
rule := &AccountStatsPricingRule{
AccountIDs: []int64{1, 2},
GroupIDs: []int64{10, 20},
}
require.False(t, matchAccountStatsRule(rule, 999, 999))
}
// ---------------------------------------------------------------------------
// findPricingForModel
// ---------------------------------------------------------------------------
func TestFindPricingForModel(t *testing.T) {
exactPricing := ChannelModelPricing{
ID: 1,
Models: []string{"claude-opus-4"},
}
wildcardPricing := ChannelModelPricing{
ID: 2,
Models: []string{"claude-*"},
}
platformPricing := ChannelModelPricing{
ID: 3,
Platform: "openai",
Models: []string{"gpt-4o"},
}
emptyPlatformPricing := ChannelModelPricing{
ID: 4,
Models: []string{"gemini-2.5-pro"},
}
tests := []struct {
name string
list []ChannelModelPricing
platform string
model string
wantID int64
wantNil bool
}{
{
name: "exact match",
list: []ChannelModelPricing{exactPricing},
platform: "anthropic",
model: "claude-opus-4",
wantID: 1,
},
{
name: "exact match case insensitive",
list: []ChannelModelPricing{{ID: 5, Models: []string{"Claude-Opus-4"}}},
platform: "",
model: "claude-opus-4",
wantID: 5,
},
{
name: "wildcard match",
list: []ChannelModelPricing{wildcardPricing},
platform: "anthropic",
model: "claude-opus-4",
wantID: 2,
},
{
name: "exact match takes priority over wildcard",
list: []ChannelModelPricing{wildcardPricing, exactPricing},
platform: "anthropic",
model: "claude-opus-4",
wantID: 1,
},
{
name: "platform mismatch skipped",
list: []ChannelModelPricing{platformPricing},
platform: "anthropic",
model: "gpt-4o",
wantNil: true,
},
{
name: "empty platform in pricing matches any",
list: []ChannelModelPricing{emptyPlatformPricing},
platform: "gemini",
model: "gemini-2.5-pro",
wantID: 4,
},
{
name: "empty platform in query matches any pricing platform",
list: []ChannelModelPricing{platformPricing},
platform: "",
model: "gpt-4o",
wantID: 3,
},
{
name: "no match at all",
list: []ChannelModelPricing{exactPricing, wildcardPricing},
platform: "anthropic",
model: "gpt-4o",
wantNil: true,
},
{
name: "empty list returns nil",
list: nil,
model: "claude-opus-4",
wantNil: true,
},
{
name: "wildcard matches by config order (first match wins)",
list: []ChannelModelPricing{
{ID: 10, Models: []string{"claude-*"}},
{ID: 11, Models: []string{"claude-opus-*"}},
},
platform: "",
model: "claude-opus-4",
wantID: 10, // config order: "claude-*" is first and matches, so it wins
},
{
name: "shorter wildcard used when longer does not match",
list: []ChannelModelPricing{
{ID: 10, Models: []string{"claude-*"}},
{ID: 11, Models: []string{"claude-opus-*"}},
},
platform: "",
model: "claude-sonnet-4",
wantID: 10, // only "claude-*" matches
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := findPricingForModel(tt.list, tt.platform, tt.model)
if tt.wantNil {
require.Nil(t, result)
return
}
require.NotNil(t, result)
require.Equal(t, tt.wantID, result.ID)
})
}
}
// ---------------------------------------------------------------------------
// calculateStatsCost
// ---------------------------------------------------------------------------
func TestCalculateStatsCost_NilPricing(t *testing.T) {
result := calculateStatsCost(nil, UsageTokens{}, 1)
require.Nil(t, result)
}
func TestCalculateStatsCost_TokenBilling(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeToken,
InputPrice: testPtrFloat64(0.001),
OutputPrice: testPtrFloat64(0.002),
}
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
}
result := calculateStatsCost(pricing, tokens, 1)
require.NotNil(t, result)
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
require.InDelta(t, 0.2, *result, 1e-12)
}
func TestCalculateStatsCost_TokenBilling_WithCache(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeToken,
InputPrice: testPtrFloat64(0.001),
OutputPrice: testPtrFloat64(0.002),
CacheWritePrice: testPtrFloat64(0.003),
CacheReadPrice: testPtrFloat64(0.0005),
}
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
CacheCreationTokens: 200,
CacheReadTokens: 300,
}
result := calculateStatsCost(pricing, tokens, 1)
require.NotNil(t, result)
// 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
// = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
require.InDelta(t, 0.95, *result, 1e-12)
}
func TestCalculateStatsCost_TokenBilling_WithImageOutput(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeToken,
InputPrice: testPtrFloat64(0.001),
OutputPrice: testPtrFloat64(0.002),
ImageOutputPrice: testPtrFloat64(0.01),
}
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
ImageOutputTokens: 10,
}
result := calculateStatsCost(pricing, tokens, 1)
require.NotNil(t, result)
// 100*0.001 + 50*0.002 + 10*0.01 = 0.1 + 0.1 + 0.1 = 0.3
require.InDelta(t, 0.3, *result, 1e-12)
}
func TestCalculateStatsCost_TokenBilling_PartialPricesNil(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeToken,
InputPrice: testPtrFloat64(0.001),
// OutputPrice, CacheWritePrice, etc. are all nil → treated as 0
}
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
CacheCreationTokens: 200,
}
result := calculateStatsCost(pricing, tokens, 1)
require.NotNil(t, result)
// Only input contributes: 100*0.001 = 0.1
require.InDelta(t, 0.1, *result, 1e-12)
}
func TestCalculateStatsCost_TokenBilling_AllTokensZero(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeToken,
InputPrice: testPtrFloat64(0.001),
OutputPrice: testPtrFloat64(0.002),
}
tokens := UsageTokens{} // all zeros
result := calculateStatsCost(pricing, tokens, 1)
// totalCost == 0 → returns nil (does not override, falls back to default formula)
require.Nil(t, result)
}
func TestCalculateStatsCost_PerRequestBilling(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModePerRequest,
PerRequestPrice: testPtrFloat64(0.05),
}
tokens := UsageTokens{InputTokens: 999, OutputTokens: 999}
result := calculateStatsCost(pricing, tokens, 3)
require.NotNil(t, result)
// 0.05 * 3 = 0.15
require.InDelta(t, 0.15, *result, 1e-12)
}
func TestCalculateStatsCost_PerRequestBilling_PriceNil(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModePerRequest,
// PerRequestPrice is nil
}
result := calculateStatsCost(pricing, UsageTokens{}, 1)
require.Nil(t, result)
}
func TestCalculateStatsCost_PerRequestBilling_PriceZero(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModePerRequest,
PerRequestPrice: testPtrFloat64(0),
}
result := calculateStatsCost(pricing, UsageTokens{}, 1)
// price == 0 → condition *pricing.PerRequestPrice > 0 is false → returns nil
require.Nil(t, result)
}
func TestCalculateStatsCost_ImageBilling(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeImage,
PerRequestPrice: testPtrFloat64(0.10),
}
result := calculateStatsCost(pricing, UsageTokens{}, 2)
require.NotNil(t, result)
// 0.10 * 2 = 0.20
require.InDelta(t, 0.20, *result, 1e-12)
}
func TestCalculateStatsCost_ImageBilling_PriceNil(t *testing.T) {
pricing := &ChannelModelPricing{
BillingMode: BillingModeImage,
// PerRequestPrice is nil
}
result := calculateStatsCost(pricing, UsageTokens{}, 1)
require.Nil(t, result)
}
func TestCalculateStatsCost_DefaultBillingMode_FallsToToken(t *testing.T) {
// BillingMode is empty string (default) → falls into token billing
pricing := &ChannelModelPricing{
InputPrice: testPtrFloat64(0.001),
OutputPrice: testPtrFloat64(0.002),
}
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
}
result := calculateStatsCost(pricing, tokens, 1)
require.NotNil(t, result)
require.InDelta(t, 0.2, *result, 1e-12)
}
// ---------------------------------------------------------------------------
// tryCustomRules — 多规则顺序测试
// ---------------------------------------------------------------------------
func TestTryCustomRules_FirstMatchWins(t *testing.T) {
channel := &Channel{
AccountStatsPricingRules: []AccountStatsPricingRule{
{
GroupIDs: []int64{1},
Pricing: []ChannelModelPricing{
{ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.01), OutputPrice: testPtrFloat64(0.02)},
},
},
{
GroupIDs: []int64{1},
Pricing: []ChannelModelPricing{
{ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.99), OutputPrice: testPtrFloat64(0.99)},
},
},
},
}
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
require.NotNil(t, result)
// 应使用第一条规则的价格:100*0.01 + 50*0.02 = 2.0
require.InDelta(t, 2.0, *result, 1e-12)
}
func TestTryCustomRules_SkipsNonMatchingRules(t *testing.T) {
channel := &Channel{
AccountStatsPricingRules: []AccountStatsPricingRule{
{
AccountIDs: []int64{888}, // 不匹配
Pricing: []ChannelModelPricing{
{ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.99)},
},
},
{
GroupIDs: []int64{1}, // 匹配
Pricing: []ChannelModelPricing{
{ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.05)},
},
},
},
}
tokens := UsageTokens{InputTokens: 100}
result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
require.NotNil(t, result)
// 跳过规则1(账号不匹配),使用规则2:100*0.05 = 5.0
require.InDelta(t, 5.0, *result, 1e-12)
}
func TestTryCustomRules_NoMatch_ReturnsNil(t *testing.T) {
channel := &Channel{
AccountStatsPricingRules: []AccountStatsPricingRule{
{
AccountIDs: []int64{888},
Pricing: []ChannelModelPricing{
{ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.01)},
},
},
},
}
tokens := UsageTokens{InputTokens: 100}
result := tryCustomRules(channel, 999, 2, "", "claude-opus-4", tokens, 1)
require.Nil(t, result) // 账号和分组都不匹配
}
func TestTryCustomRules_RuleMatchesButModelNot_ContinuesToNext(t *testing.T) {
channel := &Channel{
AccountStatsPricingRules: []AccountStatsPricingRule{
{
GroupIDs: []int64{1},
Pricing: []ChannelModelPricing{
{ID: 100, Models: []string{"gpt-4o"}, InputPrice: testPtrFloat64(0.01)}, // 模型不匹配
},
},
{
GroupIDs: []int64{1},
Pricing: []ChannelModelPricing{
{ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.05)}, // 模型匹配
},
},
},
}
tokens := UsageTokens{InputTokens: 100}
result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
require.NotNil(t, result)
require.InDelta(t, 5.0, *result, 1e-12) // 使用规则2
}
// ---------------------------------------------------------------------------
// tryModelFilePricing
// ---------------------------------------------------------------------------
// newTestBillingServiceWithPrices creates a BillingService with pre-populated
// fallback prices for testing. No config or pricing service is needed.
// The key must match what getFallbackPricing resolves to for a given model name.
// E.g., model "claude-sonnet-4" resolves to key "claude-sonnet-4".
func newTestBillingServiceWithPrices(prices map[string]*ModelPricing) *BillingService {
return &BillingService{
fallbackPrices: prices,
}
}
func TestTryModelFilePricing_Success(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
},
})
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
require.NotNil(t, result)
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
require.InDelta(t, 0.2, *result, 1e-12)
}
func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"gpt-5.6-sol": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
CacheReadPricePerToken: 0.0001,
LongContextInputThreshold: 100,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 1.5,
},
})
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "")
require.NotNil(t, result)
// Input and cache-read use the 2x input tier; output uses the 1.5x tier.
require.InDelta(t, 0.233, *result, 1e-12)
}
func TestTryModelFilePricing_AppliesServiceTierPricing(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"gpt-5.6-sol": {
InputPricePerToken: 0.001,
InputPricePerTokenPriority: 0.002,
OutputPricePerToken: 0.002,
OutputPricePerTokenPriority: 0.004,
CacheCreationPricePerToken: 0.003,
CacheCreationPricePerTokenPriority: 0.006,
CacheReadPricePerToken: 0.0005,
CacheReadPricePerTokenPriority: 0.001,
},
})
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
CacheCreationTokens: 20,
CacheReadTokens: 10,
}
tests := []struct {
name string
serviceTier string
want float64
}{
{name: "standard", serviceTier: "", want: 0.265},
{name: "priority", serviceTier: "priority", want: 0.53},
{name: "flex", serviceTier: "flex", want: 0.1325},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier)
require.NotNil(t, result)
require.InDelta(t, tt.want, *result, 1e-12)
})
}
}
func TestTryModelFilePricing_CombinesPriorityAndLongContextPricing(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"gpt-5.6-sol": {
InputPricePerToken: 0.001,
InputPricePerTokenPriority: 0.002,
OutputPricePerToken: 0.002,
OutputPricePerTokenPriority: 0.004,
CacheCreationPricePerToken: 0.003,
CacheCreationPricePerTokenPriority: 0.006,
CacheReadPricePerToken: 0.0005,
CacheReadPricePerTokenPriority: 0.001,
LongContextInputThreshold: 100,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 1.5,
},
})
tokens := UsageTokens{
InputTokens: 101,
OutputTokens: 10,
CacheCreationTokens: 5,
CacheReadTokens: 5,
}
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority")
require.NotNil(t, result)
// priority 单价先应用,再叠加长上下文输入 2x、输出 1.5x。
require.InDelta(t, 0.534, *result, 1e-12)
}
func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
// "nonexistent-model" does not match any fallback pattern
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "")
require.Nil(t, result)
}
func TestTryModelFilePricing_NilFallback(t *testing.T) {
// getFallbackPricing returns nil when key maps to nil
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": nil,
})
tokens := UsageTokens{InputTokens: 100}
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
require.Nil(t, result)
}
func TestTryModelFilePricing_ZeroCost(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
},
})
tokens := UsageTokens{} // all zero tokens → cost = 0 → nil
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
require.Nil(t, result)
}
func TestTryModelFilePricing_WithImageOutput(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
ImageOutputPricePerToken: 0.01,
},
})
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
ImageOutputTokens: 10,
}
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
require.NotNil(t, result)
// 100*0.001 + 50*0.002 + 10*0.01 = 0.1 + 0.1 + 0.1 = 0.3
require.InDelta(t, 0.3, *result, 1e-12)
}
func TestTryModelFilePricing_WithCacheTokens(t *testing.T) {
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
CacheCreationPricePerToken: 0.003,
CacheReadPricePerToken: 0.0005,
},
})
tokens := UsageTokens{
InputTokens: 100,
OutputTokens: 50,
CacheCreationTokens: 200,
CacheReadTokens: 300,
}
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
require.NotNil(t, result)
// 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
// = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
require.InDelta(t, 0.95, *result, 1e-12)
}
// ---------------------------------------------------------------------------
// resolveAccountStatsCost — integration tests covering the 4-level priority chain
// ---------------------------------------------------------------------------
func TestResolveAccountStatsCost_NilChannelService(t *testing.T) {
result := resolveAccountStatsCost(
context.Background(),
nil, // channelService is nil
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
1, 1, "claude-sonnet-4",
UsageTokens{InputTokens: 100}, 1, 0.5, "",
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_EmptyUpstreamModel(t *testing.T) {
cs := newTestChannelServiceForStats(t, &Channel{
ID: 1,
Status: StatusActive,
}, 1, "")
result := resolveAccountStatsCost(
context.Background(),
cs,
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
1, 1, "", // empty upstream model
UsageTokens{InputTokens: 100}, 1, 0.5, "",
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_GetChannelForGroupReturnsNil(t *testing.T) {
// Group 99 is NOT in the cache, so GetChannelForGroup returns nil
cs := newTestChannelServiceForStats(t, &Channel{
ID: 1,
Status: StatusActive,
}, 1, "")
result := resolveAccountStatsCost(
context.Background(),
cs,
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
1, 99, "claude-sonnet-4", // groupID 99 has no channel
UsageTokens{InputTokens: 100}, 1, 0.5, "",
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_HitsCustomRule(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
AccountStatsPricingRules: []AccountStatsPricingRule{
{
GroupIDs: []int64{10},
Pricing: []ChannelModelPricing{
{
ID: 100,
Models: []string{"claude-sonnet-4"},
InputPrice: testPtrFloat64(0.01),
OutputPrice: testPtrFloat64(0.02),
},
},
},
},
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := resolveAccountStatsCost(
context.Background(),
cs, nil, // billingService not needed when custom rule hits
1, 10, "claude-sonnet-4",
tokens, 1, 999.0, "priority", // 自定义账号价格不叠加服务层级倍率
)
require.NotNil(t, result)
// 100*0.01 + 50*0.02 = 1.0 + 1.0 = 2.0
require.InDelta(t, 2.0, *result, 1e-12)
}
func TestResolveAccountStatsCost_ApplyPricingToAccountStats_UsesTotalCost(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: true,
// No custom rules
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := resolveAccountStatsCost(
context.Background(),
cs, nil,
1, 10, "claude-sonnet-4",
tokens, 1, 0.75, "priority", // 已完成用户计费,不再重复应用服务层级倍率
)
require.NotNil(t, result)
require.InDelta(t, 0.75, *result, 1e-12)
}
func TestResolveAccountStatsCost_ApplyPricingToAccountStats_ZeroTotalCost_ReturnsNil(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: true,
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
result := resolveAccountStatsCost(
context.Background(),
cs, nil,
1, 10, "claude-sonnet-4",
UsageTokens{}, 1, 0.0, "", // totalCost = 0
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: false, // not enabled
// No custom rules
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-sonnet-4": {
InputPricePerToken: 0.001,
OutputPricePerToken: 0.002,
},
})
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := resolveAccountStatsCost(
context.Background(),
cs, bs,
1, 10, "claude-sonnet-4",
tokens, 1, 999.0, "", // totalCost ignored
)
require.NotNil(t, result)
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
require.InDelta(t, 0.2, *result, 1e-12)
}
func TestResolveAccountStatsCost_FallbackHonorsAnthropicFast(t *testing.T) {
channel := &Channel{ID: 1, Status: StatusActive}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"claude-opus-5": {
InputPricePerToken: 5e-6,
OutputPricePerToken: 25e-6,
},
})
result := resolveAccountStatsCost(
context.Background(), cs, bs,
1, 10, "claude-opus-5",
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000},
1, 0, "fast",
)
require.NotNil(t, result)
require.InDelta(t, 60, *result, 1e-12)
}
func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: false,
}
cs := newTestChannelServiceForStats(t, channel, 10, "antigravity")
bs := NewBillingService(&config.Config{}, nil)
result := resolveAccountStatsCost(
context.Background(),
cs, bs,
1, 10, "gemini-3.6-flash-low",
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "",
)
require.NotNil(t, result)
require.InDelta(t, 9.15, *result, 1e-12)
}
func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: false,
// No custom rules
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
// BillingService with no pricing for the model
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
result := resolveAccountStatsCost(
context.Background(),
cs, bs,
1, 10, "totally-unknown-model",
tokens, 1, 0.0, "",
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_NilBillingService_SkipsLiteLLM(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: false,
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
result := resolveAccountStatsCost(
context.Background(),
cs, nil, // billingService is nil
1, 10, "claude-sonnet-4",
UsageTokens{InputTokens: 100}, 1, 0.0, "",
)
require.Nil(t, result)
}
func TestResolveAccountStatsCost_CustomRulePriorityOverApplyPricing(t *testing.T) {
// Both custom rule and ApplyPricingToAccountStats are configured;
// custom rule should take precedence.
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: true,
AccountStatsPricingRules: []AccountStatsPricingRule{
{
GroupIDs: []int64{10},
Pricing: []ChannelModelPricing{
{
ID: 100,
Models: []string{"claude-sonnet-4"},
InputPrice: testPtrFloat64(0.05),
},
},
},
},
}
cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
tokens := UsageTokens{InputTokens: 100}
result := resolveAccountStatsCost(
context.Background(),
cs, nil,
1, 10, "claude-sonnet-4",
tokens, 1, 99.0, "", // totalCost = 99.0 (would be used if ApplyPricing wins)
)
require.NotNil(t, result)
// Custom rule: 100*0.05 = 5.0 (NOT 99.0 from totalCost)
require.InDelta(t, 5.0, *result, 1e-12)
}
func TestApplyAccountStatsCost_UsesUsageLogServiceTier(t *testing.T) {
channel := &Channel{
ID: 1,
Status: StatusActive,
ApplyPricingToAccountStats: false,
}
cs := newTestChannelServiceForStats(t, channel, 10, "openai")
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
"gpt-5.6-sol": {
InputPricePerToken: 0.001,
InputPricePerTokenPriority: 0.002,
OutputPricePerToken: 0.002,
OutputPricePerTokenPriority: 0.004,
},
})
serviceTier := "priority"
usageLog := &UsageLog{ServiceTier: &serviceTier}
applyAccountStatsCost(
context.Background(), usageLog, cs, bs,
1, 10, "gpt-5.6-sol", "gpt-5.6-sol",
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999,
)
require.NotNil(t, usageLog.AccountStatsCost)
require.InDelta(t, 0.4, *usageLog.AccountStatsCost, 1e-12)
}
// ---------------------------------------------------------------------------
// helpers for resolveAccountStatsCost tests
// ---------------------------------------------------------------------------
// newTestChannelServiceForStats creates a ChannelService with a single channel
// mapped to the given groupID, suitable for resolveAccountStatsCost tests.
func newTestChannelServiceForStats(t *testing.T, channel *Channel, groupID int64, platform string) *ChannelService {
t.Helper()
cache := newEmptyChannelCache()
cache.channelByGroupID[groupID] = channel
cache.groupPlatform[groupID] = platform
cs := &ChannelService{}
cache.loadedAt = time.Now()
cs.cache.Store(cache)
return cs
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,206 @@
package service
import (
"bufio"
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
"github.com/gin-gonic/gin"
)
const accountTestSuppressCompletionContextKey = "account_test_suppress_completion"
// testCNProviderAdaptiveConnection verifies every native endpoint used by an
// adaptive CN-provider account. Kimi and Zhipu use Chat Completions plus
// Anthropic; DeepSeek additionally uses its native Responses endpoint.
func (s *AccountTestService) testCNProviderAdaptiveConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
testModelID := strings.TrimSpace(modelID)
if testModelID == "" {
testModelID = openai.DefaultTestModel
}
testModelID = account.GetMappedModel(testModelID)
authToken := strings.TrimSpace(account.GetOpenAIProtocolAPIKey())
if authToken == "" {
return s.sendErrorAndEnd(c, "No API key available")
}
// The existing Chat probe owns the SSE lifecycle. Suppress intermediate
// completion events until every native adaptive endpoint has passed.
c.Set(accountTestSuppressCompletionContextKey, true)
defer c.Set(accountTestSuppressCompletionContextKey, false)
if err := s.testCNProviderChatCompletionsConnection(c, account, modelID, prompt); err != nil {
return err
}
if err := s.testCNProviderAdaptiveAnthropicConnection(c, account, testModelID, authToken); err != nil {
return err
}
if account.Platform == PlatformDeepseek {
if err := s.testCNProviderAdaptiveResponsesConnection(c, account, testModelID, authToken); err != nil {
return err
}
}
c.Set(accountTestSuppressCompletionContextKey, false)
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
return nil
}
func (s *AccountTestService) testCNProviderAdaptiveAnthropicConnection(c *gin.Context, account *Account, testModelID string, authToken string) error {
ctx := c.Request.Context()
baseURL, err := s.validateUpstreamBaseURL(account.GetCNProtocolBaseURL(APIProtocolAnthropic))
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid adaptive Anthropic base URL: %s", err.Error()))
}
apiURL := strings.TrimRight(baseURL, "/") + "/v1/messages"
payload, err := createTestPayload(testModelID)
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create adaptive Anthropic test payload")
}
payloadBytes, _ := json.Marshal(payload)
s.sendEvent(c, TestEvent{Type: "status", Text: "正在通过原生 /v1/messages 测试自适应 Anthropic 端点"})
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create adaptive Anthropic request")
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("anthropic-version", "2023-06-01")
for key, value := range claude.DefaultHeaders {
req.Header.Set(key, value)
}
req.Header.Set("anthropic-beta", claude.APIKeyBetaHeader)
setAnthropicAPIKeyAuthHeader(req.Header, account, authToken)
account.ApplyHeaderOverrides(req.Header)
resp, err := s.doCNProviderAdaptiveRequest(req, account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Adaptive Anthropic endpoint request failed: %s", err.Error()))
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
errMsg := fmt.Sprintf("Adaptive Anthropic endpoint returned %d: %s", resp.StatusCode, string(body))
if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil {
_ = s.accountRepo.SetError(ctx, account.ID, errMsg)
}
return s.sendErrorAndEnd(c, errMsg)
}
if err := s.processCNProviderAdaptiveAnthropicStream(c, resp.Body); err != nil {
return err
}
s.sendEvent(c, TestEvent{Type: "status", Text: "已通过原生 /v1/messages 验证"})
return nil
}
func (s *AccountTestService) processCNProviderAdaptiveAnthropicStream(c *gin.Context, body io.Reader) error {
reader := bufio.NewReader(body)
for {
line, err := reader.ReadString('\n')
if err != nil {
if err == io.EOF {
return s.sendErrorAndEnd(c, "Adaptive Anthropic stream ended before message_stop")
}
return s.sendErrorAndEnd(c, fmt.Sprintf("Adaptive Anthropic stream read error: %s", err.Error()))
}
line = strings.TrimSpace(line)
if line == "" || !sseDataPrefix.MatchString(line) {
continue
}
jsonStr := sseDataPrefix.ReplaceAllString(line, "")
if jsonStr == "[DONE]" {
return nil
}
var data map[string]any
if err := json.Unmarshal([]byte(jsonStr), &data); err != nil {
continue
}
switch eventType, _ := data["type"].(string); eventType {
case "content_block_delta":
if delta, ok := data["delta"].(map[string]any); ok {
if text, ok := delta["text"].(string); ok && text != "" {
s.sendEvent(c, TestEvent{Type: "content", Text: text})
}
}
case "message_stop":
return nil
case "error":
errorMsg := "Unknown error"
if errData, ok := data["error"].(map[string]any); ok {
if message, ok := errData["message"].(string); ok && message != "" {
errorMsg = message
}
}
return s.sendErrorAndEnd(c, fmt.Sprintf("Adaptive Anthropic endpoint error: %s", errorMsg))
}
}
}
func (s *AccountTestService) testCNProviderAdaptiveResponsesConnection(c *gin.Context, account *Account, testModelID string, authToken string) error {
ctx := c.Request.Context()
baseURL, err := s.validateUpstreamBaseURL(account.GetCNProtocolBaseURL(APIProtocolResponses))
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid adaptive Responses base URL: %s", err.Error()))
}
apiURL := buildOpenAIResponsesURLForPlatform(account.Platform, baseURL)
payload := createOpenAITestPayload(testModelID, false)
// DeepSeek's native Responses endpoint is stateless and does not need the
// OpenAI probe's synthetic instructions.
delete(payload, "instructions")
payloadBytes, _ := json.Marshal(payload)
payloadBytes = normalizeDeepSeekResponsesRequestBody(account, payloadBytes)
s.sendEvent(c, TestEvent{Type: "status", Text: "正在通过原生 /responses 测试自适应 Responses 端点"})
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create adaptive Responses request")
}
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("Authorization", "Bearer "+authToken)
applyOpenAICodexProbeHeaders(req.Header)
account.ApplyHeaderOverrides(req.Header)
resp, err := s.doCNProviderAdaptiveRequest(req, account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Adaptive Responses endpoint request failed: %s", err.Error()))
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
errMsg := fmt.Sprintf("Adaptive Responses endpoint returned %d: %s", resp.StatusCode, string(body))
if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil {
_ = s.accountRepo.SetError(ctx, account.ID, errMsg)
}
return s.sendErrorAndEnd(c, errMsg)
}
if err := s.processOpenAIStream(c, resp.Body); err != nil {
return err
}
s.sendEvent(c, TestEvent{Type: "status", Text: "已通过原生 /responses 验证"})
return nil
}
func (s *AccountTestService) doCNProviderAdaptiveRequest(req *http.Request, account *Account) (*http.Response, error) {
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
return s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
}
@@ -0,0 +1,191 @@
//go:build unit
package service
import (
"io"
"net/http"
"strings"
"testing"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func adaptiveCNAccountTestAccount(id int64, platform string) *Account {
return &Account{
ID: id,
Name: "adaptive-cn-test",
Platform: platform,
Type: AccountTypeAPIKey,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-adaptive-test",
"api_protocol": APIProtocolAdaptive,
"api_base_urls": map[string]any{
APIProtocolChatCompletions: "http://chat.example/v1",
APIProtocolAnthropic: "http://anthropic.example",
APIProtocolResponses: "http://responses.example",
},
},
}
}
func adaptiveCNAccountTestService(account *Account, responses ...*http.Response) (*AccountTestService, *httpUpstreamRecorder) {
repo := &openAIAccountTestRepo{
mockAccountRepoForGemini: mockAccountRepoForGemini{
accountsByID: map[int64]*Account{account.ID: account},
},
}
upstream := &httpUpstreamRecorder{responses: responses}
return &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
cfg: rawChatCompletionsTestConfig(),
}, upstream
}
func adaptiveCNChatTestResponse() *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(`data: {"choices":[{"delta":{"content":"chat ok"},"finish_reason":"stop"}]}
data: [DONE]
`)),
}
}
func adaptiveCNAnthropicTestResponse() *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(`data: {"type":"content_block_delta","delta":{"text":"anthropic ok"}}
data: {"type":"message_stop"}
`)),
}
}
func adaptiveCNResponsesTestResponse() *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(`data: {"type":"response.output_text.delta","delta":"responses ok"}
data: {"type":"response.completed"}
`)),
}
}
func TestAccountTestService_AdaptiveChatOnlyProvidersTestChatAndAnthropicEndpoints(t *testing.T) {
for index, testCase := range []struct {
name string
platform string
model string
}{
{name: "Kimi", platform: PlatformKimi, model: "kimi-k2.5"},
{name: "Zhipu", platform: PlatformZhipu, model: "glm-4.7"},
} {
t.Run(testCase.name, func(t *testing.T) {
account := adaptiveCNAccountTestAccount(int64(301+index), testCase.platform)
svc, upstream := adaptiveCNAccountTestService(
account,
adaptiveCNChatTestResponse(),
adaptiveCNAnthropicTestResponse(),
)
c, recorder := newTestContext()
err := svc.TestAccountConnection(c, account.ID, testCase.model, "hello", AccountTestModeDefault)
require.NoError(t, err)
require.Len(t, upstream.requests, 2)
require.Equal(t, "http://chat.example/v1/chat/completions", upstream.requests[0].URL.String())
require.Equal(t, "http://anthropic.example/v1/messages", upstream.requests[1].URL.String())
require.Equal(t, "Bearer sk-adaptive-test", upstream.requests[0].Header.Get("Authorization"))
require.Equal(t, "sk-adaptive-test", upstream.requests[1].Header.Get("x-api-key"))
require.Equal(t, 1, strings.Count(recorder.Body.String(), `"type":"test_start"`))
require.Equal(t, 1, strings.Count(recorder.Body.String(), `"type":"test_complete"`))
require.Contains(t, recorder.Body.String(), "已通过原生 /v1/messages 验证")
})
}
}
func TestAccountTestService_AdaptiveDeepSeekAlsoTestsResponsesEndpoint(t *testing.T) {
account := adaptiveCNAccountTestAccount(302, PlatformDeepseek)
svc, upstream := adaptiveCNAccountTestService(
account,
adaptiveCNChatTestResponse(),
adaptiveCNAnthropicTestResponse(),
adaptiveCNResponsesTestResponse(),
)
c, recorder := newTestContext()
err := svc.TestAccountConnection(c, account.ID, "deepseek-chat", "", AccountTestModeDefault)
require.NoError(t, err)
require.Len(t, upstream.requests, 3)
require.Equal(t, "http://responses.example/responses", upstream.requests[2].URL.String())
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.requests[2].Context()))
require.Equal(t, "Bearer sk-adaptive-test", upstream.requests[2].Header.Get("Authorization"))
require.True(t, gjson.GetBytes(upstream.bodies[2], "stream").Bool())
require.False(t, gjson.GetBytes(upstream.bodies[2], "store").Bool())
require.False(t, gjson.GetBytes(upstream.bodies[2], "instructions").Exists())
require.Equal(t, 1, strings.Count(recorder.Body.String(), `"type":"test_complete"`))
require.Contains(t, recorder.Body.String(), "已通过原生 /responses 验证")
}
func TestAccountTestService_AdaptiveStopsAndNamesFailingEndpoint(t *testing.T) {
account := adaptiveCNAccountTestAccount(303, PlatformDeepseek)
svc, upstream := adaptiveCNAccountTestService(
account,
adaptiveCNChatTestResponse(),
newJSONResponse(http.StatusNotFound, `{"error":{"message":"missing messages route"}}`),
)
c, recorder := newTestContext()
err := svc.TestAccountConnection(c, account.ID, "deepseek-chat", "", AccountTestModeDefault)
require.Error(t, err)
require.Contains(t, err.Error(), "Adaptive Anthropic endpoint returned 404")
require.Len(t, upstream.requests, 2)
require.Contains(t, recorder.Body.String(), `"type":"error"`)
require.NotContains(t, recorder.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_AdaptiveRejectsInvalidAnthropicSuccessBody(t *testing.T) {
account := adaptiveCNAccountTestAccount(305, PlatformKimi)
svc, upstream := adaptiveCNAccountTestService(
account,
adaptiveCNChatTestResponse(),
newJSONResponse(http.StatusOK, `<html>not an Anthropic stream</html>`),
)
c, recorder := newTestContext()
err := svc.TestAccountConnection(c, account.ID, "kimi-k2.5", "", AccountTestModeDefault)
require.Error(t, err)
require.Contains(t, err.Error(), "Adaptive Anthropic stream ended before message_stop")
require.Len(t, upstream.requests, 2)
require.Contains(t, recorder.Body.String(), `"type":"error"`)
require.NotContains(t, recorder.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_FixedCNChatProtocolStillTestsOnlyChatEndpoint(t *testing.T) {
account := adaptiveCNAccountTestAccount(304, PlatformZhipu)
account.Credentials["api_protocol"] = APIProtocolChatCompletions
account.Credentials["base_url"] = "http://fixed-chat.example/v1"
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNChatTestResponse())
c, recorder := newTestContext()
err := svc.TestAccountConnection(c, account.ID, "glm-4.7", "", AccountTestModeDefault)
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
require.Equal(t, "http://fixed-chat.example/v1/chat/completions", upstream.requests[0].URL.String())
require.Equal(t, 1, strings.Count(recorder.Body.String(), `"type":"test_complete"`))
}
@@ -0,0 +1,59 @@
//go:build unit
package service
import (
"encoding/json"
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestCreateGeminiTestPayload_ImageModel(t *testing.T) {
t.Parallel()
payload := createGeminiTestPayload("gemini-2.5-flash-image", "draw a tiny robot")
var parsed struct {
Contents []struct {
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
} `json:"contents"`
GenerationConfig struct {
ResponseModalities []string `json:"responseModalities"`
ImageConfig struct {
AspectRatio string `json:"aspectRatio"`
} `json:"imageConfig"`
} `json:"generationConfig"`
}
require.NoError(t, json.Unmarshal(payload, &parsed))
require.Len(t, parsed.Contents, 1)
require.Len(t, parsed.Contents[0].Parts, 1)
require.Equal(t, "draw a tiny robot", parsed.Contents[0].Parts[0].Text)
require.Equal(t, []string{"TEXT", "IMAGE"}, parsed.GenerationConfig.ResponseModalities)
require.Equal(t, "1:1", parsed.GenerationConfig.ImageConfig.AspectRatio)
}
func TestProcessGeminiStream_EmitsImageEvent(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
svc := &AccountTestService{}
stream := strings.NewReader("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"ok\"},{\"inlineData\":{\"mimeType\":\"image/png\",\"data\":\"QUJD\"}}]}}]}\n\ndata: [DONE]\n\n")
err := svc.processGeminiStream(ctx, stream)
require.NoError(t, err)
body := recorder.Body.String()
require.Contains(t, body, "\"type\":\"content\"")
require.Contains(t, body, "\"text\":\"ok\"")
require.Contains(t, body, "\"type\":\"image\"")
require.Contains(t, body, "\"image_url\":\"data:image/png;base64,QUJD\"")
require.Contains(t, body, "\"mime_type\":\"image/png\"")
}
@@ -0,0 +1,644 @@
//go:build unit
package service
import (
"bytes"
"context"
"encoding/base64"
"errors"
"image"
"image/color"
"image/png"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
type grokAccountTestRateLimitRepo struct {
*mockAccountRepoForGemini
rateLimitedCalls int
resetAt time.Time
}
func TestObserveGrokTestResponseClassifiesBodyOnlyQuotaErrors(t *testing.T) {
account := &Account{ID: 1901, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
svc := &AccountTestService{accountRepo: repo}
resp := &http.Response{
StatusCode: http.StatusBadRequest,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"error":{"code":"subscription:free-usage-exhausted","message":"included free usage exhausted"}}`)),
}
svc.observeGrokTestResponse(context.Background(), account, resp)
require.Equal(t, 1, repo.tempUnschedCalls)
require.Equal(t, "grok free usage exhausted", repo.lastTempUnschedReason)
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
require.Contains(t, string(body), "free-usage-exhausted")
}
func TestObserveGrokTestResponseDoesNotQuarantineContentPolicy(t *testing.T) {
account := &Account{ID: 1902, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
svc := &AccountTestService{accountRepo: repo}
resp := &http.Response{
StatusCode: http.StatusForbidden,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`)),
}
svc.observeGrokTestResponse(context.Background(), account, resp)
require.Zero(t, repo.tempUnschedCalls)
require.Zero(t, repo.rateLimitedCalls)
}
func TestObserveGrokTestResponseKeepsEntitlement403Cooldown(t *testing.T) {
account := &Account{ID: 1903, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
svc := &AccountTestService{accountRepo: repo}
resp := &http.Response{
StatusCode: http.StatusForbidden,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"subscription required"}}`)),
}
before := time.Now()
svc.observeGrokTestResponse(context.Background(), account, resp)
require.Equal(t, 1, repo.tempUnschedCalls)
require.Equal(t, "grok entitlement or subscription tier denied", repo.lastTempUnschedReason)
require.Greater(t, repo.lastTempUnschedUntil, before.Add(29*time.Minute))
}
func (r *grokAccountTestRateLimitRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
r.rateLimitedCalls++
r.resetAt = resetAt
return nil
}
func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 13,
Name: "grok-oauth",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
"model_mapping": map[string]any{
"grok": "grok-4.3",
},
},
}
repo := &mockAccountRepoForGemini{
accountsByID: map[int64]*Account{account.ID: account},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
"data: {\"type\":\"response.completed\"}\n\n",
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/13/test", nil)
err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
require.NoError(t, err)
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer grok-access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "application/json, text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, grokQuotaProbeInput, gjson.GetBytes(upstream.lastBody, "input").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "max_output_tokens").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists())
require.NotContains(t, rec.Body.String(), "claude")
require.Contains(t, rec.Body.String(), `"model":"grok-4.3"`)
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 16,
Name: "grok-oauth-default-model",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
"data: {\"type\":\"response.completed\"}\n\n",
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/16/test", nil)
err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeDefault)
require.NoError(t, err)
require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String())
require.Contains(t, recorder.Body.String(), `"model":"grok-4.5"`)
}
func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 14,
Name: "grok-oauth-limited",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{"Retry-After": []string{"45"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/14/test", nil)
err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
require.Error(t, err)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, time.Now().Add(45*time.Second), repo.resetAt, time.Second)
}
func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"quota exhausted"}}`)),
}}
svc := &AccountTestService{
accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, nil), httpUpstream: upstream,
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/15/test", nil)
before := time.Now()
err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
require.Error(t, err)
require.Equal(t, 1, repo.rateLimitedCalls)
require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.resetAt, time.Second)
}
func TestAccountTestService_GrokImageModelUsesImagesGenerations(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 17, Name: "grok-oauth-image", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
"base_url": "https://cli-chat-proxy.grok.com/v1",
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"data":[{"b64_json":"QUJD","mime_type":"image/jpeg"}]}`,
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/17/test", nil)
err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image", "a red apple", AccountTestModeDefault)
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/images/generations", upstream.lastReq.URL.String())
require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "a red apple", gjson.GetBytes(upstream.lastBody, "prompt").String())
require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
require.Contains(t, rec.Body.String(), `"type":"image"`)
require.Contains(t, rec.Body.String(), "data:image/jpeg;base64,QUJD")
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_GrokWebSearchModeUsesResponsesWebSearchTool(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 18, Name: "grok-oauth-search", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(
`{"id":"r1","output":[{"type":"web_search_call","id":"ws1","status":"completed"},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Grok is built by xAI."}]}]}`,
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/18/test", nil)
err := svc.TestAccountConnection(c, account.ID, "grok-4.5", "xAI Grok", AccountTestModeGrokSearch)
require.NoError(t, err)
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
// Standalone web_search wraps the query in the gateway-style prompt.
require.Contains(t, gjson.GetBytes(upstream.lastBody, "input").String(), "xAI Grok")
require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
require.Equal(t, "web_search_call.action.sources", gjson.GetBytes(upstream.lastBody, "include.0").String())
require.Contains(t, rec.Body.String(), "web_search ok")
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_GrokTTSIncludesLanguage(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 19, Name: "grok-oauth-tts", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"audio/mpeg"}},
Body: io.NopCloser(strings.NewReader("ID3fakeaudio")),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/19/test", nil)
err := svc.TestAccountConnection(c, account.ID, "", "hello voice", AccountTestModeGrokTTS)
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/tts", upstream.lastReq.URL.String())
require.Equal(t, "hello voice", gjson.GetBytes(upstream.lastBody, "text").String())
require.Equal(t, "en", gjson.GetBytes(upstream.lastBody, "language").String())
require.Contains(t, rec.Body.String(), "tts ok")
require.Contains(t, rec.Body.String(), `"type":"audio"`)
require.Contains(t, rec.Body.String(), "data:audio/mpeg;base64,")
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_GrokImageEditUsesUploadedImage(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 24, Name: "grok-oauth-image-edit", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"data":[{"b64_json":"QUJD","mime_type":"image/png"}]}`,
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/24/test", nil)
// 8x8 solid PNG (xAI min dimension is 8px).
src := minimalAccountTestPNGDataURL(8, 8)
err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image", "edit me", AccountTestModeGrokImage, AccountTestOptions{
ImageDataURL: src,
})
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/images/edits", upstream.lastReq.URL.String())
require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.url").String(), "data:image/png;base64,"))
require.Equal(t, "image_url", gjson.GetBytes(upstream.lastBody, "image.type").String())
require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
// concrete image model ids pass through; only bare "grok-imagine" is aliased.
require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
require.Contains(t, rec.Body.String(), `"type":"image"`)
}
func TestAccountTestService_GrokImageEditRejectsTinySource(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 25, Name: "grok-oauth-image-tiny", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: &httpUpstreamRecorder{},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/25/test", nil)
// 1x1 PNG data URL
tiny := "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image-quality", "edit", AccountTestModeGrokImage, AccountTestOptions{
ImageDataURL: tiny,
})
require.Error(t, err)
require.Contains(t, rec.Body.String(), "too small")
}
// minimalAccountTestPNGDataURL builds a solid RGBA PNG as a data URL for tests.
func minimalAccountTestPNGDataURL(w, h int) string {
img := image.NewRGBA(image.Rect(0, 0, w, h))
for y := 0; y < h; y++ {
for x := 0; x < w; x++ {
img.Set(x, y, color.RGBA{R: 200, G: 40, B: 40, A: 255})
}
}
var buf bytes.Buffer
_ = png.Encode(&buf, img)
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())
}
func TestAccountTestService_GrokExplicitImageModeDefaultsModel(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 20, Name: "grok-oauth-image-mode", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"data":[{"b64_json":"QUJD","mime_type":"image/jpeg"}]}`,
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/20/test", nil)
err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokImage)
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/images/generations", upstream.lastReq.URL.String())
require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_GrokVideoUpstreamErrorIsNotMaskedAsSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 21, Name: "grok-oauth-video-err", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"code":"invalid-argument","error":"bad video request"}`,
)),
}}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/21/test", nil)
err := svc.TestAccountConnection(c, account.ID, "grok-imagine-video", "bounce ball", AccountTestModeGrokVideo)
require.Error(t, err)
require.Equal(t, "https://api.x.ai/v1/videos/generations", upstream.lastReq.URL.String())
require.Contains(t, rec.Body.String(), `"type":"error"`)
require.Contains(t, rec.Body.String(), "Grok videos API returned 400")
require.NotContains(t, rec.Body.String(), `"success":true`)
}
type grokRealtimeTestConn struct {
msg []byte
}
func (c *grokRealtimeTestConn) WriteJSON(context.Context, any) error { return nil }
func (c *grokRealtimeTestConn) ReadMessage(context.Context) ([]byte, error) {
if c == nil || len(c.msg) == 0 {
return nil, context.DeadlineExceeded
}
return c.msg, nil
}
func (c *grokRealtimeTestConn) Ping(context.Context) error { return nil }
func (c *grokRealtimeTestConn) Close() error { return nil }
type grokRealtimeTestDialer struct {
lastURL string
lastAuth string
lastProxy string
conn openAIWSClientConn
err error
status int
}
func (d *grokRealtimeTestDialer) Dial(_ context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error) {
d.lastURL = wsURL
d.lastAuth = headers.Get("Authorization")
d.lastProxy = proxyURL
if d.err != nil {
return nil, d.status, nil, d.err
}
if d.conn == nil {
d.conn = &grokRealtimeTestConn{}
}
return d.conn, 0, nil, nil
}
func TestAccountTestService_GrokRealtimeModeDialsWS(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 22, Name: "grok-oauth-realtime", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
dialer := &grokRealtimeTestDialer{
conn: &grokRealtimeTestConn{msg: []byte(`{"type":"session.created","session":{"id":"sess_1"}}`)},
}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
grokWSDialer: dialer,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", nil)
err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokRealtime)
require.NoError(t, err)
require.Contains(t, dialer.lastURL, "wss://api.x.ai/v1/realtime")
require.Contains(t, dialer.lastURL, "model=grok-voice-latest")
require.Equal(t, "Bearer grok-access-token", dialer.lastAuth)
require.Contains(t, rec.Body.String(), "realtime ws handshake ok")
require.Contains(t, rec.Body.String(), "session.created")
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
require.Contains(t, rec.Body.String(), `"success":true`)
}
func TestAccountTestService_GrokRealtimeModeDialFailure(t *testing.T) {
gin.SetMode(gin.TestMode)
account := &Account{
ID: 23, Name: "grok-oauth-realtime-fail", Platform: PlatformGrok,
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
Credentials: map[string]any{
"access_token": "grok-access-token",
"refresh_token": "grok-refresh-token",
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
},
}
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
dialer := &grokRealtimeTestDialer{
status: 401,
err: &openAIWSHandshakeError{Body: []byte(`{"error":"unauthorized"}`), Err: errors.New("websocket handshake failed")},
}
svc := &AccountTestService{
accountRepo: repo,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
grokWSDialer: dialer,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/23/test", nil)
err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokRealtime)
require.Error(t, err)
require.Contains(t, rec.Body.String(), `"type":"error"`)
require.Contains(t, rec.Body.String(), "Realtime")
}
@@ -0,0 +1,330 @@
package service
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
// compactProbeSSESuccessBody 是原生 v2 压缩成功的最小 SSE 形态:
// output_item.done 携带 compaction item + response.completed。
const compactProbeSSESuccessBody = "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"id\":\"cmp_probe\",\"encrypted_content\":\"blob\"}}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_probe\",\"output\":[]}}\n\n"
func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersistsSupport(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 1,
Name: "openai-oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
"chatgpt_account_is_fedramp": true,
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid-probe"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/1/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
// 原生 v2:探测普通 /responses 线,不再打已下线的 /responses/compact。
require.Equal(t, chatgptCodexAPIURL, upstream.lastReq.URL.String())
require.Equal(t, "chatgpt.com", upstream.lastReq.Host)
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
require.NotEmpty(t, upstream.lastReq.Header.Get("Session_Id"))
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
require.Equal(t, codexCLIUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "chatgpt-acc", upstream.lastReq.Header.Get("chatgpt-account-id"))
require.Equal(t, "true", upstream.lastReq.Header.Get("x-openai-fedramp"))
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool())
inputItems := gjson.GetBytes(upstream.lastBody, "input").Array()
require.NotEmpty(t, inputItems)
require.Equal(t, "compaction_trigger", inputItems[len(inputItems)-1].Get("type").String())
updates := <-updateCalls
require.Equal(t, true, updates["openai_compact_supported"])
require.Equal(t, http.StatusOK, updates["openai_compact_last_status"])
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
}
func TestAccountTestService_TestAccountConnection_OpenAICompactOAuth404MarksUnsupported(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 2,
Name: "openai-oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusNotFound,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`404 page not found`)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/2/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.Error(t, err)
updates := <-updateCalls
require.Equal(t, false, updates["openai_compact_supported"])
require.Equal(t, http.StatusNotFound, updates["openai_compact_last_status"])
require.Contains(t, rec.Body.String(), `"type":"error"`)
}
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesNativeResponsesPath(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 3,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com/v1",
// post-#5641compact_model_mapping 仅作用于 legacy /responses/compact
// 原生 v2 探测不应用它。
"compact_model_mapping": map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"},
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/3/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
require.Equal(t, "https://example.com/v1/responses", upstream.lastReq.URL.String())
requireOpenAICodexProbeHeaders(t, upstream.lastReq.Header)
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String(),
"原生 v2 探测不应用 compact_model_mapping")
updates := <-updateCalls
require.Equal(t, true, updates["openai_compact_supported"])
}
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBaseURLUsesResponsesPath(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 4,
Name: "openai-apikey-default",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/4/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String())
<-updateCalls
}
func TestAccountTestService_TestAccountConnection_OpenAICompact2xxWithoutItemMarksUnsupported(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 5,
Name: "openai-oauth-no-item",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
// 200 但流里没有 compaction item:链路吞掉了 compaction_trigger 的形态
//#5478 的 "got 0 items"),必须判定为不支持。
noItemBody := "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"message\",\"id\":\"msg_1\"}}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_x\",\"output\":[]}}\n\n"
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(noItemBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/5/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.Error(t, err)
updates := <-updateCalls
require.Equal(t, false, updates["openai_compact_supported"])
require.Contains(t, rec.Body.String(), `"type":"error"`)
}
// 探测与真实转发走同一 /responses 端点,出站身份必须与真实 Codex 同构:
// session/thread 为 UUID、携带 x-codex-installation-id(收敛账号用收敛值)。
func TestAccountTestService_TestAccountConnection_OpenAICompactProbeIdentityMatchesRealTraffic(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 6,
Name: "openai-oauth-identity",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
// 收敛是显式 opt-in(#5610),这里显式开启以验证探测身份与真实流量同构。
Extra: map[string]any{
"codex_fingerprint_mode": "session",
codexFingerprintSeedExtraKey: testCodexFingerprintSeed,
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/6/test", bytes.NewReader(nil))
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
// 显式 session 收敛模式:出站身份 = 账号级收敛值
seed, ok := codexFingerprintSeed(account.Extra)
require.True(t, ok)
converged := resolveConvergedSessionID(seed)
require.Equal(t, converged, upstream.lastReq.Header.Get("session-id"))
require.Equal(t, converged, upstream.lastReq.Header.Get("session_id"))
require.Equal(t, resolveConvergedInstallationID(&account, seed), upstream.lastReq.Header.Get("x-codex-installation-id"),
"真实 Codex 每个请求必带 installation-id,探测不得缺失")
require.NotContains(t, upstream.lastReq.Header.Get("session-id"), "probe_compact",
"探测标识不得是可被上游一眼识别的字面量")
<-updateCalls
}
func TestCompactProbeSessionID_IsUUIDShaped(t *testing.T) {
for _, id := range []int64{0, 1, 987654} {
got := compactProbeSessionID(id)
_, err := uuid.Parse(got)
require.NoError(t, err, "探测会话标识必须是 UUID 形态: %s", got)
}
require.Equal(t, compactProbeSessionID(7), compactProbeSessionID(7), "同账号应稳定复用同一会话")
require.NotEqual(t, compactProbeSessionID(7), compactProbeSessionID(8))
}
@@ -0,0 +1,93 @@
package service
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAccountTestService_OpenAIImageOAuthHandlesOutputItemDoneFallback(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/1/test", nil)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
},
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"aGVsbG8=\",\"revised_prompt\":\"draw a cat\",\"output_format\":\"png\"}}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000006,\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[]}}\n\n" +
"data: [DONE]\n\n",
)),
},
}
svc := &AccountTestService{httpUpstream: upstream}
account := &Account{
ID: 53,
Name: "openai-oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"access_token": "token-123",
},
}
err := svc.testOpenAIImageOAuth(c, context.Background(), account, "gpt-image-2", "draw a cat")
require.NoError(t, err)
require.NotNil(t, upstream.lastReq)
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
require.Contains(t, rec.Body.String(), "Calling Codex /responses image tool")
require.Contains(t, rec.Body.String(), "data:image/png;base64,aGVsbG8=")
require.Contains(t, rec.Body.String(), "\"success\":true")
}
func TestAccountTestService_OpenAIImageAPIKeyUsesConfiguredV1BaseURL(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/1/test", nil)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
},
Body: io.NopCloser(strings.NewReader(`{"data":[{"b64_json":"aGVsbG8=","revised_prompt":"draw a cat"}]}`)),
},
}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{},
}
account := &Account{
ID: 54,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"api_key": "test-api-key",
"base_url": "https://image-upstream.example/v1",
},
}
err := svc.testOpenAIImageAPIKey(c, context.Background(), account, "gpt-image-2", "draw a cat")
require.NoError(t, err)
require.NotNil(t, upstream.lastReq)
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
require.Equal(t, "https://image-upstream.example/v1/images/generations", upstream.lastReq.URL.String())
require.Equal(t, "Bearer test-api-key", upstream.lastReq.Header.Get("Authorization"))
require.Contains(t, rec.Body.String(), "data:image/png;base64,aGVsbG8=")
require.Contains(t, rec.Body.String(), "\"success\":true")
}
@@ -0,0 +1,599 @@
//go:build unit
package service
import (
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
// --- shared test helpers ---
type queuedHTTPUpstream struct {
responses []*http.Response
requests []*http.Request
tlsFlags []bool
}
func (u *queuedHTTPUpstream) Do(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
return nil, fmt.Errorf("unexpected Do call")
}
func (u *queuedHTTPUpstream) DoWithTLS(req *http.Request, _ string, _ int64, _ int, profile *tlsfingerprint.Profile) (*http.Response, error) {
u.requests = append(u.requests, req)
u.tlsFlags = append(u.tlsFlags, profile != nil)
if len(u.responses) == 0 {
return nil, fmt.Errorf("no mocked response")
}
resp := u.responses[0]
u.responses = u.responses[1:]
return resp, nil
}
func newJSONResponse(status int, body string) *http.Response {
return &http.Response{
StatusCode: status,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(body)),
}
}
// --- test functions ---
func newTestContext() (*gin.Context, *httptest.ResponseRecorder) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/1/test", nil)
return c, rec
}
type openAIAccountTestRepo struct {
mockAccountRepoForGemini
updatedExtra map[string]any
bulkUpdatedIDs []int64
bulkUpdatedPayload AccountBulkUpdate
rateLimitedID int64
rateLimitedAt *time.Time
clearedErrorID int64
setErrorID int64
setErrorMsg string
}
func (r *openAIAccountTestRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
r.updatedExtra = updates
return nil
}
func (r *openAIAccountTestRepo) BulkUpdate(_ context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
r.bulkUpdatedIDs = append([]int64(nil), ids...)
r.bulkUpdatedPayload = updates
return int64(len(ids)), nil
}
func (r *openAIAccountTestRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error {
r.rateLimitedID = id
r.rateLimitedAt = &resetAt
return nil
}
func (r *openAIAccountTestRepo) ClearError(_ context.Context, id int64) error {
r.clearedErrorID = id
return nil
}
func (r *openAIAccountTestRepo) SetError(_ context.Context, id int64, errorMsg string) error {
r.setErrorID = id
r.setErrorMsg = errorMsg
return nil
}
func TestAccountTestService_OpenAISuccessPersistsSnapshotFromHeaders(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.completed"}
`))
resp.Header.Set("x-codex-primary-used-percent", "88")
resp.Header.Set("x-codex-primary-reset-after-seconds", "604800")
resp.Header.Set("x-codex-primary-window-minutes", "10080")
resp.Header.Set("x-codex-secondary-used-percent", "42")
resp.Header.Set("x-codex-secondary-reset-after-seconds", "18000")
resp.Header.Set("x-codex-secondary-window-minutes", "300")
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 89,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.requests[0].Context()))
require.NotEmpty(t, repo.updatedExtra)
require.Equal(t, 42.0, repo.updatedExtra["codex_5h_used_percent"])
require.Equal(t, 88.0, repo.updatedExtra["codex_7d_used_percent"])
require.Contains(t, recorder.Body.String(), "test_complete")
}
func TestAccountTestService_OpenAIOAuthTestNormalizesGPT56Alias(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.completed"}
`))
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{httpUpstream: upstream}
account := &Account{
ID: 90,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.6", "", "")
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
body, err := io.ReadAll(upstream.requests[0].Body)
require.NoError(t, err)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(body, "model").String())
}
func TestAccountTestService_OpenAIShadowUsesParentCredentialsAndShadowModel(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.completed"}
`))
parentID := int64(100)
parent := &Account{
ID: parentID,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"access_token": "parent-token",
"chatgpt_account_id": "org-parent",
},
}
shadow := &Account{
ID: 200,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
ParentAccountID: &parentID,
QuotaDimension: QuotaDimensionSpark,
Concurrency: 2,
Credentials: map[string]any{
"model_mapping": map[string]any{
"gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
},
},
}
repo := &openAIAccountTestRepo{
mockAccountRepoForGemini: mockAccountRepoForGemini{
accountsByID: map[int64]*Account{
parentID: parent,
200: shadow,
},
},
}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
err := svc.TestAccountConnection(ctx, shadow.ID, "gpt-5.3-codex-spark", "", "")
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
req := upstream.requests[0]
require.Equal(t, "Bearer parent-token", req.Header.Get("Authorization"))
require.Equal(t, "org-parent", req.Header.Get("chatgpt-account-id"))
body, err := io.ReadAll(req.Body)
require.NoError(t, err)
require.Equal(t, "gpt-5.3-codex-spark", gjson.GetBytes(body, "model").String())
require.Contains(t, recorder.Body.String(), `"success":true`)
}
func TestAccountTestService_OpenAIStreamEOFBeforeCompletedFails(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.output_text.delta","delta":"hi"}
`))
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{httpUpstream: upstream}
account := &Account{
ID: 90,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Contains(t, recorder.Body.String(), "response.completed")
require.NotContains(t, recorder.Body.String(), `"success":true`)
}
func TestAccountTestService_OpenAI429PersistsSnapshotAndRateLimitState(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusTooManyRequests, `{"error":{"type":"usage_limit_reached","message":"limit reached","resets_at":1777283883}}`)
resp.Header.Set("x-codex-primary-used-percent", "100")
resp.Header.Set("x-codex-primary-reset-after-seconds", "604800")
resp.Header.Set("x-codex-primary-window-minutes", "10080")
resp.Header.Set("x-codex-secondary-used-percent", "100")
resp.Header.Set("x-codex-secondary-reset-after-seconds", "18000")
resp.Header.Set("x-codex-secondary-window-minutes", "300")
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 88,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusError,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.NotEmpty(t, repo.updatedExtra)
require.Equal(t, 100.0, repo.updatedExtra["codex_5h_used_percent"])
require.Equal(t, account.ID, repo.rateLimitedID)
require.NotNil(t, repo.rateLimitedAt)
require.Equal(t, account.ID, repo.clearedErrorID)
require.Equal(t, StatusActive, account.Status)
require.Empty(t, account.ErrorMessage)
require.NotNil(t, account.RateLimitResetAt)
}
func TestAccountTestService_OpenAI429BodyOnlyPersistsRateLimitAndClearsStaleError(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusTooManyRequests, `{"error":{"type":"usage_limit_reached","message":"limit reached","resets_at":"1777283883"}}`)
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 77,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusError,
ErrorMessage: "Access forbidden (403): account may be suspended or lack permissions",
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, account.ID, repo.rateLimitedID)
require.NotNil(t, repo.rateLimitedAt)
require.Equal(t, account.ID, repo.clearedErrorID)
require.Equal(t, StatusActive, account.Status)
require.Empty(t, account.ErrorMessage)
require.NotNil(t, account.RateLimitResetAt)
require.Empty(t, repo.updatedExtra)
}
func TestAccountTestService_OpenAI429SyncsObservedPlanType(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusTooManyRequests, `{"error":{"type":"usage_limit_reached","message":"limit reached","plan_type":"free","resets_at":1777283883}}`)
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 81,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token", "plan_type": "plus"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, []int64{account.ID}, repo.bulkUpdatedIDs)
require.Equal(t, "free", repo.bulkUpdatedPayload.Credentials["plan_type"])
require.Equal(t, "free", account.Credentials["plan_type"])
require.Equal(t, account.ID, repo.rateLimitedID)
require.NotNil(t, account.RateLimitResetAt)
}
func TestAccountTestService_OpenAI429ActiveAccountDoesNotClearError(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusTooManyRequests, `{"error":{"type":"usage_limit_reached","message":"limit reached","resets_in_seconds":3600}}`)
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 78,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, account.ID, repo.rateLimitedID)
require.NotNil(t, repo.rateLimitedAt)
require.Zero(t, repo.clearedErrorID)
require.Equal(t, StatusActive, account.Status)
require.NotNil(t, account.RateLimitResetAt)
}
func TestAccountTestService_OpenAI429WithoutResetSignalDoesNotMutateRuntimeState(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusTooManyRequests, `{"error":{"type":"usage_limit_reached","message":"limit reached"}}`)
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 79,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusError,
ErrorMessage: "stale 403",
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Zero(t, repo.rateLimitedID)
require.Nil(t, repo.rateLimitedAt)
require.Zero(t, repo.clearedErrorID)
require.Equal(t, StatusError, account.Status)
require.Equal(t, "stale 403", account.ErrorMessage)
require.Nil(t, account.RateLimitResetAt)
}
func TestAccountTestService_OpenAI401SetsPermanentErrorOnly(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusUnauthorized, `{"error":"bad token"}`)
repo := &openAIAccountTestRepo{}
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
account := &Account{
ID: 80,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, account.ID, repo.setErrorID)
require.Contains(t, repo.setErrorMsg, "Authentication failed (401)")
require.Zero(t, repo.rateLimitedID)
require.Zero(t, repo.clearedErrorID)
require.Nil(t, account.RateLimitResetAt)
}
func TestAccountTestService_OpenAIAPIKeyResponsesUsesCodexProbeHeaders(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader("data: {\"type\":\"response.completed\"}\n\n"))
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 95,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://compat-upstream.example/v1",
},
Extra: map[string]any{openai_compat.ExtraKeyResponsesSupported: true},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
req := upstream.requests[0]
require.Equal(t, "https://compat-upstream.example/v1/responses", req.URL.String())
requireOpenAICodexProbeHeaders(t, req.Header)
}
func TestAccountTestService_OpenAIAPIKeyResponsesUnsupportedUsesChatCompletionsPath(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
upstreamBody := strings.Join([]string{
`data: {"id":"chatcmpl_test","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"pong"},"finish_reason":null}]}`,
"",
`data: {"id":"chatcmpl_test","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`,
"",
"data: [DONE]",
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamBody)),
}}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 91,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://compat-upstream.example/v1",
},
Extra: map[string]any{openai_compat.ExtraKeyResponsesSupported: false},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "hello", "")
require.NoError(t, err)
require.NotNil(t, upstream.lastReq)
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
require.Equal(t, "https://compat-upstream.example/v1/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, "Bearer sk-test", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.Equal(t, "hello", gjson.GetBytes(upstream.lastBody, "messages.0.content").String())
require.False(t, gjson.GetBytes(upstream.lastBody, "input").Exists())
body := recorder.Body.String()
require.Contains(t, body, "pong")
require.Contains(t, body, "已通过 /v1/chat/completions 验证")
require.Contains(t, body, `"success":true`)
require.NotContains(t, body, "当前测试接口仅支持 Responses API 路径")
}
func TestAccountTestService_OpenAIChatCompletionsPathReturns4xx(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
upstream := &httpUpstreamRecorder{resp: newJSONResponse(http.StatusBadRequest, `{"error":{"message":"bad request"}}`)}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 92,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://compat-upstream.example",
},
Extra: map[string]any{openai_compat.ExtraKeyResponsesSupported: false},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, "https://compat-upstream.example/v1/chat/completions", upstream.lastReq.URL.String())
require.Contains(t, err.Error(), "Chat Completions API (/v1/chat/completions) returned 400")
require.Contains(t, recorder.Body.String(), "/v1/chat/completions")
require.NotContains(t, recorder.Body.String(), `"success":true`)
}
func TestAccountTestService_OpenAIChatCompletionsPathTimeout(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
upstream := &httpUpstreamRecorder{err: context.DeadlineExceeded}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 93,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://compat-upstream.example",
},
Extra: map[string]any{openai_compat.ExtraKeyResponsesSupported: false},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, "https://compat-upstream.example/v1/chat/completions", upstream.lastReq.URL.String())
require.Contains(t, err.Error(), "Chat Completions API (/v1/chat/completions) request failed")
require.Contains(t, err.Error(), context.DeadlineExceeded.Error())
require.Contains(t, recorder.Body.String(), "/v1/chat/completions")
require.NotContains(t, recorder.Body.String(), `"success":true`)
}
func TestAccountTestService_OpenAIChatCompletionsPathRejectsNonJSONStream(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader("data: not-json\n\n")),
}}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 94,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://compat-upstream.example",
},
Extra: map[string]any{openai_compat.ExtraKeyResponsesSupported: false},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
require.Error(t, err)
require.Equal(t, "https://compat-upstream.example/v1/chat/completions", upstream.lastReq.URL.String())
require.Contains(t, err.Error(), "Invalid Chat Completions response from /v1/chat/completions")
require.Contains(t, recorder.Body.String(), "/v1/chat/completions")
require.NotContains(t, recorder.Body.String(), `"success":true`)
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,191 @@
package service
import (
"context"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
)
// Minimal UsageLogRepository stub for batch usage tests (HEAD lacks geminiUsageLogRepoStub).
type usageBatchLogRepoStub struct{}
var _ UsageLogRepository = (*usageBatchLogRepoStub)(nil)
func (r *usageBatchLogRepoStub) Create(context.Context, *UsageLog) (bool, error) {
return false, nil
}
func (r *usageBatchLogRepoStub) GetByID(context.Context, int64) (*UsageLog, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) Delete(context.Context, int64) error { return nil }
func (r *usageBatchLogRepoStub) ListByUser(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByAPIKey(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByAccount(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByUserAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByAPIKeyAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByAccountAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) ListByModelAndTimeRange(context.Context, string, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) GetAccountWindowStats(context.Context, int64, time.Time) (*usagestats.AccountStats, error) {
return &usagestats.AccountStats{}, nil
}
func (r *usageBatchLogRepoStub) GetAccountTodayStats(context.Context, int64) (*usagestats.AccountStats, error) {
return &usagestats.AccountStats{}, nil
}
func (r *usageBatchLogRepoStub) GetDashboardStats(context.Context) (*usagestats.DashboardStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUsageTrendWithFilters(context.Context, time.Time, time.Time, string, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.TrendDataPoint, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetModelStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, *int16, *bool, *int8) ([]usagestats.ModelStat, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetEndpointStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.EndpointStat, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUpstreamEndpointStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.EndpointStat, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetGroupStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, *int16, *bool, *int8) ([]usagestats.GroupStat, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserBreakdownStats(context.Context, time.Time, time.Time, usagestats.UserBreakdownDimension, int) ([]usagestats.UserBreakdownItem, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAllGroupUsageSummary(context.Context, time.Time) ([]usagestats.GroupUsageSummary, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAPIKeyUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.APIKeyUsageTrendPoint, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.UserUsageTrendPoint, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserSpendingRanking(context.Context, time.Time, time.Time, int) (*usagestats.UserSpendingRankingResponse, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetBatchUserUsageStats(context.Context, []int64, time.Time, time.Time) (map[int64]*usagestats.BatchUserUsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetBatchAPIKeyUsageStats(context.Context, []int64, time.Time, time.Time) (map[int64]*usagestats.BatchAPIKeyUsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserDashboardStats(context.Context, int64) (*usagestats.UserDashboardStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAPIKeyDashboardStats(context.Context, int64) (*usagestats.UserDashboardStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserUsageTrendByUserID(context.Context, int64, time.Time, time.Time, string) ([]usagestats.TrendDataPoint, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserModelStats(context.Context, int64, time.Time, time.Time) ([]usagestats.ModelStat, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) ListWithFilters(context.Context, pagination.PaginationParams, usagestats.UsageLogFilters) ([]UsageLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *usageBatchLogRepoStub) GetGlobalStats(context.Context, time.Time, time.Time) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetStatsWithFilters(context.Context, usagestats.UsageLogFilters) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAccountUsageStats(context.Context, int64, time.Time, time.Time) (*usagestats.AccountUsageStatsResponse, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetUserStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAPIKeyStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetAccountStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetModelStatsAggregated(context.Context, string, time.Time, time.Time) (*usagestats.UsageStats, error) {
return nil, nil
}
func (r *usageBatchLogRepoStub) GetDailyStatsAggregated(context.Context, int64, time.Time, time.Time) ([]map[string]any, error) {
return nil, nil
}
func TestAccountUsageService_GetUsageBatch_BestEffortByAccount(t *testing.T) {
t.Parallel()
resetAt := time.Now().Add(2 * time.Hour).UTC().Truncate(time.Second)
repo := &stubOpenAIAccountRepo{
accounts: []Account{
{
ID: 7001,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{
"passive_usage_7d_utilization": 0.62,
},
},
{
ID: 7002,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_usage_updated_at": time.Now().UTC().Format(time.RFC3339),
"codex_5h_used_percent": 18.0,
"codex_5h_reset_at": resetAt.Format(time.RFC3339),
"codex_7d_used_percent": 34.0,
"codex_7d_reset_at": resetAt.Add(24 * time.Hour).Format(time.RFC3339),
"workspace_id": "org-test",
"chatgpt_account_id": "acct-test",
"openai_snapshot_version": "test",
},
},
{
ID: 7003,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
},
},
}
svc := &AccountUsageService{
accountRepo: repo,
usageLogRepo: &usageBatchLogRepoStub{},
cache: NewUsageCache(),
}
usageByAccount, errorsByAccount, err := svc.GetUsageBatch(context.Background(), []int64{7001, 7002, 7003, 7002}, false)
if err != nil {
t.Fatalf("GetUsageBatch() error = %v", err)
}
if usageByAccount[7001] == nil || usageByAccount[7001].Source != "passive" {
t.Fatalf("expected anthropic passive usage, got %#v", usageByAccount[7001])
}
if usageByAccount[7002] == nil || usageByAccount[7002].FiveHour == nil || usageByAccount[7002].FiveHour.Utilization != 18.0 {
t.Fatalf("expected openai snapshot usage, got %#v", usageByAccount[7002])
}
if !strings.Contains(strings.ToLower(errorsByAccount[7003]), "does not support usage query") {
t.Fatalf("expected API key account error to be preserved, got %q", errorsByAccount[7003])
}
}
@@ -0,0 +1,119 @@
package service
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestClaudeUsageResponse_FableWindowDecoding(t *testing.T) {
t.Run("seven_day_overage_included", func(t *testing.T) {
raw := `{
"five_hour": {"utilization": 12.0, "resets_at": "2026-07-03T10:00:00Z"},
"seven_day": {"utilization": 34.0, "resets_at": "2026-07-08T00:00:00Z"},
"seven_day_overage_included": {"utilization": 56.0, "resets_at": "2026-07-08T03:00:00Z"}
}`
var resp ClaudeUsageResponse
require.NoError(t, json.Unmarshal([]byte(raw), &resp))
require.Equal(t, 56.0, resp.SevenDayOverageIncluded.Utilization)
require.Equal(t, "2026-07-08T03:00:00Z", resp.SevenDayOverageIncluded.ResetsAt)
})
t.Run("absent", func(t *testing.T) {
raw := `{"five_hour": {"utilization": 12.0, "resets_at": "2026-07-03T10:00:00Z"}}`
var resp ClaudeUsageResponse
require.NoError(t, json.Unmarshal([]byte(raw), &resp))
require.Zero(t, resp.SevenDayOverageIncluded.Utilization)
require.Empty(t, resp.SevenDayOverageIncluded.ResetsAt)
})
}
func TestBuildUsageInfo_SevenDayFable(t *testing.T) {
svc := &AccountUsageService{}
now := time.Now()
resetAt := now.Add(72 * time.Hour).UTC().Truncate(time.Second)
var resp ClaudeUsageResponse
resp.FiveHour.Utilization = 10
resp.SevenDayOverageIncluded = ClaudeUsageWindow{
Utilization: 88,
ResetsAt: resetAt.Format(time.RFC3339),
}
info := svc.buildUsageInfo(&resp, &now)
require.NotNil(t, info.SevenDayFable)
require.Equal(t, 88.0, info.SevenDayFable.Utilization)
require.NotNil(t, info.SevenDayFable.ResetsAt)
require.True(t, info.SevenDayFable.ResetsAt.Equal(resetAt))
require.Greater(t, info.SevenDayFable.RemainingSeconds, 0)
// 无 Fable 数据时不应创建窗口
var empty ClaudeUsageResponse
empty.FiveHour.Utilization = 10
info = svc.buildUsageInfo(&empty, &now)
require.Nil(t, info.SevenDayFable)
}
func TestBuildPassiveUsageWindow(t *testing.T) {
future := time.Now().Add(48 * time.Hour).Unix()
t.Run("utilization and reset", func(t *testing.T) {
window := buildPassiveUsageWindow(map[string]any{
"passive_usage_7d_oi_utilization": 0.87,
"passive_usage_7d_oi_reset": float64(future),
}, "passive_usage_7d_oi_utilization", "passive_usage_7d_oi_reset")
require.NotNil(t, window)
require.InDelta(t, 87.0, window.Utilization, 1e-9)
require.NotNil(t, window.ResetsAt)
require.Equal(t, future, window.ResetsAt.Unix())
require.Greater(t, window.RemainingSeconds, 0)
})
t.Run("no data returns nil", func(t *testing.T) {
require.Nil(t, buildPassiveUsageWindow(nil, "u", "r"))
require.Nil(t, buildPassiveUsageWindow(map[string]any{}, "u", "r"))
})
t.Run("expired reset clamps remaining to zero", func(t *testing.T) {
past := time.Now().Add(-time.Hour).Unix()
window := buildPassiveUsageWindow(map[string]any{
"u": 0.5,
"r": float64(past),
}, "u", "r")
require.NotNil(t, window)
require.Equal(t, 0, window.RemainingSeconds)
})
t.Run("utilization only", func(t *testing.T) {
window := buildPassiveUsageWindow(map[string]any{"u": 0.25}, "u", "r")
require.NotNil(t, window)
require.InDelta(t, 25.0, window.Utilization, 1e-9)
require.Nil(t, window.ResetsAt)
})
}
func TestSyncActiveToPassive_WritesFableExtras(t *testing.T) {
repo := &accountUsageCodexProbeRepo{updateExtraCh: make(chan map[string]any, 1)}
svc := &AccountUsageService{accountRepo: repo}
resetAt := time.Now().Add(72 * time.Hour).Truncate(time.Second)
usage := &UsageInfo{
SevenDayFable: &UsageProgress{
Utilization: 87,
ResetsAt: &resetAt,
},
}
svc.syncActiveToPassive(t.Context(), 1, usage)
select {
case updates := <-repo.updateExtraCh:
require.InDelta(t, 0.87, updates["passive_usage_7d_oi_utilization"], 1e-9)
require.Equal(t, resetAt.Unix(), updates["passive_usage_7d_oi_reset"])
require.Contains(t, updates, "passive_usage_sampled_at")
default:
t.Fatal("expected UpdateExtra to be called with fable extras")
}
}
@@ -0,0 +1,151 @@
package service
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// sparkShadowUsageTestRepo is a minimal AccountRepository stub for spark shadow
// usage tests. GetByID serves both shadow and parent accounts from a map;
// UpdateExtra records the persisted updates for assertion.
type sparkShadowUsageTestRepo struct {
AccountRepository
accounts map[int64]*Account
updateExtraCh chan map[string]any
}
func (r *sparkShadowUsageTestRepo) GetByID(_ context.Context, id int64) (*Account, error) {
if acc, ok := r.accounts[id]; ok {
return acc, nil
}
return nil, fmt.Errorf("account %d not found", id)
}
func (r *sparkShadowUsageTestRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
if r.updateExtraCh != nil {
copied := make(map[string]any, len(updates))
for k, v := range updates {
copied[k] = v
}
r.updateExtraCh <- copied
}
return nil
}
// TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows covers
// two assertions required by Task 3.2:
//
// A) After getOpenAIUsage on a spark shadow account the shadow row's
// Extra["codex_5h_used_percent"] is persisted, and the upstream call carried
// the PARENT account's chatgpt-account-id (not the shadow's empty one).
//
// B) (P1-b regression guard) The UsageInfo RETURNED by the same call has
// non-nil FiveHour AND SevenDay windows — proving that the rebuild happened
// and not just the DB write.
func TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows(t *testing.T) {
t.Parallel()
ctx := context.Background()
pid := int64(100)
shadow := &Account{
ID: 200,
ParentAccountID: &pid,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
QuotaDimension: QuotaDimensionSpark,
}
parent := &Account{
ID: 100,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"chatgpt_account_id": "org-spark-parent",
},
}
// Repo shared by both the OpenAIQuotaService (needs shadow+parent for resolve)
// and the AccountUsageService (needs UpdateExtra for persist).
updateExtraCh := make(chan map[string]any, 1)
repo := &sparkShadowUsageTestRepo{
accounts: map[int64]*Account{200: shadow, 100: parent},
updateExtraCh: updateExtraCh,
}
// Token cache: return a fake token for the parent account key.
tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
OpenAITokenCacheKey(parent): "fake-access-token",
}}
tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
// httptest server: records the chatgpt-account-id header and returns a
// synthetic OpenAIQuotaUsage with codex_bengalfox 5h+7d windows.
var capturedAccountID string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capturedAccountID = r.Header.Get("chatgpt-account-id")
w.Header().Set("content-type", "application/json")
resp := OpenAIQuotaUsage{
AdditionalRateLimits: []OpenAIAdditionalRateLimit{
{
MeteredFeature: "codex_bengalfox",
RateLimit: &OpenAIRateLimit{
// Primary window → 5h (18000 s = 300 min)
PrimaryWindow: &OpenAIRateLimitWindow{
UsedPercent: 42.5,
ResetAfterSeconds: 3600,
LimitWindowSeconds: 18000,
},
// Secondary window → 7d (604800 s = 10080 min)
SecondaryWindow: &OpenAIRateLimitWindow{
UsedPercent: 10.0,
ResetAfterSeconds: 86400,
LimitWindowSeconds: 604800,
},
},
},
},
}
_ = json.NewEncoder(w).Encode(resp)
}))
defer srv.Close()
quotaService := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
svc := &AccountUsageService{
accountRepo: repo,
openAIQuotaService: quotaService,
}
usage, err := svc.getOpenAIUsage(ctx, shadow, true /*force*/)
require.NoError(t, err)
// Assertion A-1: upstream received the PARENT's chatgpt-account-id.
require.Equal(t, "org-spark-parent", capturedAccountID,
"QueryUsage must use parent's chatgpt-account-id for spark shadow accounts")
// Assertion A-2: shadow Extra was persisted with codex_5h_used_percent.
select {
case updates := <-updateExtraCh:
require.Contains(t, updates, "codex_5h_used_percent",
"persisted extra must contain codex_5h_used_percent")
require.InDelta(t, 42.5, updates["codex_5h_used_percent"], 0.01,
"codex_5h_used_percent must match the upstream value")
case <-time.After(2 * time.Second):
t.Fatal("UpdateExtra was not called within timeout — spark shadow persist did not happen")
}
// Assertion B (P1-b regression guard): returned UsageInfo must have
// non-nil windows. This FAILS if the code only writes Extra without
// rebuilding the returned UsageInfo.
require.NotNil(t, usage.FiveHour,
"returned UsageInfo.FiveHour must be non-nil (rebuild from merged Extra must happen)")
require.NotNil(t, usage.SevenDay,
"returned UsageInfo.SevenDay must be non-nil (rebuild from merged Extra must happen)")
}
@@ -0,0 +1,258 @@
package service
import (
"context"
"net/http"
"testing"
"time"
)
type accountUsageCodexProbeRepo struct {
stubOpenAIAccountRepo
updateExtraCh chan map[string]any
rateLimitCh chan time.Time
}
func (r *accountUsageCodexProbeRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
if r.updateExtraCh != nil {
copied := make(map[string]any, len(updates))
for k, v := range updates {
copied[k] = v
}
r.updateExtraCh <- copied
}
return nil
}
func (r *accountUsageCodexProbeRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
if r.rateLimitCh != nil {
r.rateLimitCh <- resetAt
}
return nil
}
func TestShouldRefreshOpenAICodexSnapshot(t *testing.T) {
t.Parallel()
rateLimitedUntil := time.Now().Add(5 * time.Minute)
now := time.Now()
usage := &UsageInfo{
FiveHour: &UsageProgress{Utilization: 0},
SevenDay: &UsageProgress{Utilization: 0},
}
if !shouldRefreshOpenAICodexSnapshot(&Account{RateLimitResetAt: &rateLimitedUntil}, usage, now) {
t.Fatal("expected rate-limited account to force codex snapshot refresh")
}
if shouldRefreshOpenAICodexSnapshot(&Account{}, usage, now) {
t.Fatal("expected complete non-rate-limited usage to skip codex snapshot refresh")
}
if !shouldRefreshOpenAICodexSnapshot(&Account{}, &UsageInfo{FiveHour: nil, SevenDay: &UsageProgress{}}, now) {
t.Fatal("expected missing 5h snapshot to require refresh")
}
staleAt := now.Add(-(openAIProbeCacheTTL + time.Minute)).Format(time.RFC3339)
if !shouldRefreshOpenAICodexSnapshot(&Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"openai_oauth_responses_websockets_v2_enabled": true,
"codex_usage_updated_at": staleAt,
},
}, usage, now) {
t.Fatal("expected stale ws snapshot to trigger refresh")
}
}
// TestShouldRefreshOpenAICodexSnapshot_SparkShadowIgnoresWSv2 外审第9轮 P1:spark 影子用量走
// QueryUsage(/wham/usage,与 WSv2 无关),staleness 不得被 WSv2 门控,否则首刷后窗口永久冻结。
func TestShouldRefreshOpenAICodexSnapshot_SparkShadowIgnoresWSv2(t *testing.T) {
t.Parallel()
now := time.Now()
usage := &UsageInfo{
FiveHour: &UsageProgress{Utilization: 0},
SevenDay: &UsageProgress{Utilization: 0},
}
staleAt := now.Add(-(openAIProbeCacheTTL + time.Minute)).Format(time.RFC3339)
freshAt := now.Add(-time.Minute).Format(time.RFC3339)
parentID := int64(7001)
// 影子无 WSv2,但首刷后窗口已存在;过期 codex_usage_updated_at 必须触发再刷新。
shadowStale := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
ParentAccountID: &parentID,
QuotaDimension: QuotaDimensionSpark,
Extra: map[string]any{"codex_usage_updated_at": staleAt},
}
if !shouldRefreshOpenAICodexSnapshot(shadowStale, usage, now) {
t.Fatal("expected stale spark shadow (no WSv2) to trigger refresh")
}
// 影子时间戳仍新鲜→不刷(TTL 生效)。
shadowFresh := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
ParentAccountID: &parentID,
QuotaDimension: QuotaDimensionSpark,
Extra: map[string]any{"codex_usage_updated_at": freshAt},
}
if shouldRefreshOpenAICodexSnapshot(shadowFresh, usage, now) {
t.Fatal("expected fresh spark shadow to skip refresh (TTL not elapsed)")
}
// 反向对照:普通账号无 WSv2 + 过期时间戳→仍不刷(WSv2 门控普通账号的 probe 刷新)。
normalNoWS := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{"codex_usage_updated_at": staleAt},
}
if shouldRefreshOpenAICodexSnapshot(normalNoWS, usage, now) {
t.Fatal("expected non-WSv2 normal account to skip codex probe refresh")
}
}
func TestExtractOpenAICodexProbeUpdatesAccepts429WithCodexHeaders(t *testing.T) {
t.Parallel()
headers := make(http.Header)
headers.Set("x-codex-primary-used-percent", "100")
headers.Set("x-codex-primary-reset-after-seconds", "604800")
headers.Set("x-codex-primary-window-minutes", "10080")
headers.Set("x-codex-secondary-used-percent", "100")
headers.Set("x-codex-secondary-reset-after-seconds", "18000")
headers.Set("x-codex-secondary-window-minutes", "300")
updates, err := extractOpenAICodexProbeUpdates(&http.Response{StatusCode: http.StatusTooManyRequests, Header: headers})
if err != nil {
t.Fatalf("extractOpenAICodexProbeUpdates() error = %v", err)
}
if len(updates) == 0 {
t.Fatal("expected codex probe updates from 429 headers")
}
if got := updates["codex_5h_used_percent"]; got != 100.0 {
t.Fatalf("codex_5h_used_percent = %v, want 100", got)
}
if got := updates["codex_7d_used_percent"]; got != 100.0 {
t.Fatalf("codex_7d_used_percent = %v, want 100", got)
}
}
func TestAccountUsageService_PersistOpenAICodexProbeSnapshotOnlyUpdatesExtra(t *testing.T) {
t.Parallel()
repo := &accountUsageCodexProbeRepo{
updateExtraCh: make(chan map[string]any, 1),
rateLimitCh: make(chan time.Time, 1),
}
svc := &AccountUsageService{accountRepo: repo}
svc.persistOpenAICodexProbeSnapshot(321, map[string]any{
"codex_7d_used_percent": 100.0,
"codex_7d_reset_at": time.Now().Add(2 * time.Hour).UTC().Truncate(time.Second).Format(time.RFC3339),
})
select {
case updates := <-repo.updateExtraCh:
if got := updates["codex_7d_used_percent"]; got != 100.0 {
t.Fatalf("codex_7d_used_percent = %v, want 100", got)
}
case <-time.After(2 * time.Second):
t.Fatal("等待 codex 探测快照写入 extra 超时")
}
select {
case got := <-repo.rateLimitCh:
t.Fatalf("不应将探测快照写入运行时限流状态: %v", got)
case <-time.After(200 * time.Millisecond):
}
}
func TestAccountUsageService_GetOpenAIUsage_DoesNotPromoteCodexExtraToRateLimit(t *testing.T) {
t.Parallel()
resetAt := time.Now().Add(6 * 24 * time.Hour).UTC().Truncate(time.Second)
repo := &accountUsageCodexProbeRepo{
rateLimitCh: make(chan time.Time, 1),
}
svc := &AccountUsageService{accountRepo: repo}
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_5h_used_percent": 1.0,
"codex_5h_reset_at": time.Now().Add(2 * time.Hour).UTC().Truncate(time.Second).Format(time.RFC3339),
"codex_7d_used_percent": 100.0,
"codex_7d_reset_at": resetAt.Format(time.RFC3339),
},
}
usage, err := svc.getOpenAIUsage(context.Background(), account, false)
if err != nil {
t.Fatalf("getOpenAIUsage() error = %v", err)
}
if usage.SevenDay == nil || usage.SevenDay.Utilization != 100.0 {
t.Fatalf("预期 7 天用量仍然可见,实际为 %#v", usage.SevenDay)
}
if account.RateLimitResetAt != nil {
t.Fatalf("不应让已耗尽的 codex extra 改写运行时限流状态: %v", account.RateLimitResetAt)
}
select {
case got := <-repo.rateLimitCh:
t.Fatalf("不应将已耗尽的 codex extra 持久化为运行时限流状态: %v", got)
case <-time.After(200 * time.Millisecond):
}
}
func TestBuildCodexUsageProgressFromExtra_ZerosExpiredWindow(t *testing.T) {
t.Parallel()
now := time.Date(2026, 3, 16, 12, 0, 0, 0, time.UTC)
t.Run("expired 5h window zeroes utilization", func(t *testing.T) {
extra := map[string]any{
"codex_5h_used_percent": 42.0,
"codex_5h_reset_at": "2026-03-16T10:00:00Z", // 2h ago
}
progress := buildCodexUsageProgressFromExtra(extra, "5h", now)
if progress == nil {
t.Fatal("expected non-nil progress")
}
if progress.Utilization != 0 {
t.Fatalf("expected Utilization=0 for expired window, got %v", progress.Utilization)
}
if progress.RemainingSeconds != 0 {
t.Fatalf("expected RemainingSeconds=0, got %v", progress.RemainingSeconds)
}
})
t.Run("active 5h window keeps utilization", func(t *testing.T) {
resetAt := now.Add(2 * time.Hour).Format(time.RFC3339)
extra := map[string]any{
"codex_5h_used_percent": 42.0,
"codex_5h_reset_at": resetAt,
}
progress := buildCodexUsageProgressFromExtra(extra, "5h", now)
if progress == nil {
t.Fatal("expected non-nil progress")
}
if progress.Utilization != 42.0 {
t.Fatalf("expected Utilization=42, got %v", progress.Utilization)
}
})
t.Run("expired 7d window zeroes utilization", func(t *testing.T) {
extra := map[string]any{
"codex_7d_used_percent": 88.0,
"codex_7d_reset_at": "2026-03-15T00:00:00Z", // yesterday
}
progress := buildCodexUsageProgressFromExtra(extra, "7d", now)
if progress == nil {
t.Fatal("expected non-nil progress")
}
if progress.Utilization != 0 {
t.Fatalf("expected Utilization=0 for expired 7d window, got %v", progress.Utilization)
}
})
}
@@ -0,0 +1,135 @@
package service
import (
"context"
"sync"
"testing"
"time"
)
// sessionWindowSyncRepo 记录 syncActiveToPassive 触发的所有写操作。
type sessionWindowSyncRepo struct {
AccountRepository
mu sync.Mutex
extraUpdates []map[string]any
sessionWindowEnds []sessionWindowEndCall
}
type sessionWindowEndCall struct {
AccountID int64
End time.Time
}
func (r *sessionWindowSyncRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
copied := make(map[string]any, len(updates))
for k, v := range updates {
copied[k] = v
}
r.extraUpdates = append(r.extraUpdates, copied)
return nil
}
func (r *sessionWindowSyncRepo) UpdateSessionWindowEnd(_ context.Context, id int64, end time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
r.sessionWindowEnds = append(r.sessionWindowEnds, sessionWindowEndCall{AccountID: id, End: end})
return nil
}
func TestEstimateSetupTokenUsage_ExpiredWindowZeroes(t *testing.T) {
t.Parallel()
past := time.Now().Add(-2 * time.Hour)
svc := &AccountUsageService{}
info := svc.estimateSetupTokenUsage(&Account{
SessionWindowEnd: &past,
Extra: map[string]any{
"session_window_utilization": 0.53,
},
})
if info.FiveHour == nil {
t.Fatal("expected non-nil FiveHour info")
}
if info.FiveHour.Utilization != 0 {
t.Fatalf("expected Utilization=0 for expired window, got %v", info.FiveHour.Utilization)
}
if info.FiveHour.ResetsAt != nil {
t.Fatalf("expected ResetsAt=nil for expired window, got %v", info.FiveHour.ResetsAt)
}
if info.FiveHour.RemainingSeconds != 0 {
t.Fatalf("expected RemainingSeconds=0 for expired window, got %v", info.FiveHour.RemainingSeconds)
}
}
func TestEstimateSetupTokenUsage_ActiveWindowPreservesUtilization(t *testing.T) {
t.Parallel()
future := time.Now().Add(3 * time.Hour)
svc := &AccountUsageService{}
info := svc.estimateSetupTokenUsage(&Account{
SessionWindowEnd: &future,
Extra: map[string]any{
"session_window_utilization": 0.53,
},
})
if info.FiveHour == nil {
t.Fatal("expected non-nil FiveHour info")
}
if info.FiveHour.Utilization != 53 {
t.Fatalf("expected Utilization=53, got %v", info.FiveHour.Utilization)
}
if info.FiveHour.ResetsAt == nil || !info.FiveHour.ResetsAt.Equal(future) {
t.Fatalf("expected ResetsAt=%v, got %v", future, info.FiveHour.ResetsAt)
}
if info.FiveHour.RemainingSeconds <= 0 {
t.Fatalf("expected positive RemainingSeconds, got %v", info.FiveHour.RemainingSeconds)
}
}
func TestSyncActiveToPassive_WritesFiveHourSessionWindowEnd(t *testing.T) {
t.Parallel()
repo := &sessionWindowSyncRepo{}
svc := &AccountUsageService{accountRepo: repo}
resetsAt := time.Now().Add(3 * time.Hour).UTC().Truncate(time.Second)
svc.syncActiveToPassive(context.Background(), 42, &UsageInfo{
FiveHour: &UsageProgress{
Utilization: 53,
ResetsAt: &resetsAt,
},
})
repo.mu.Lock()
defer repo.mu.Unlock()
if len(repo.sessionWindowEnds) != 1 {
t.Fatalf("expected 1 UpdateSessionWindowEnd call, got %d", len(repo.sessionWindowEnds))
}
call := repo.sessionWindowEnds[0]
if call.AccountID != 42 {
t.Fatalf("expected AccountID=42, got %d", call.AccountID)
}
if !call.End.Equal(resetsAt) {
t.Fatalf("expected End=%v, got %v", resetsAt, call.End)
}
}
func TestSyncActiveToPassive_SkipsSessionWindowEndWhenResetMissing(t *testing.T) {
t.Parallel()
repo := &sessionWindowSyncRepo{}
svc := &AccountUsageService{accountRepo: repo}
svc.syncActiveToPassive(context.Background(), 99, &UsageInfo{
FiveHour: &UsageProgress{Utilization: 10},
})
repo.mu.Lock()
defer repo.mu.Unlock()
if len(repo.sessionWindowEnds) != 0 {
t.Fatalf("expected no UpdateSessionWindowEnd calls when ResetsAt is nil, got %d", len(repo.sessionWindowEnds))
}
}
@@ -0,0 +1,105 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestGetWebSearchEmulationMode_Enabled(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: "enabled"},
}
require.Equal(t, WebSearchModeEnabled, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_Disabled(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: "disabled"},
}
require.Equal(t, WebSearchModeDisabled, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_Default(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: "default"},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_UnknownString(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: "unknown"},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_OldBoolTrue(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: true},
}
// bool true → tolerant fallback → enabled (not default)
require.Equal(t, WebSearchModeEnabled, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_OldBoolFalse(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: false},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_NilAccount(t *testing.T) {
var a *Account
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_NilExtra(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: nil,
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_MissingField(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Extra: map[string]any{},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_NonAnthropicPlatform(t *testing.T) {
a := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Extra: map[string]any{featureKeyWebSearchEmulation: "enabled"},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
func TestGetWebSearchEmulationMode_NonAPIKeyType(t *testing.T) {
a := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Extra: map[string]any{featureKeyWebSearchEmulation: "enabled"},
}
require.Equal(t, WebSearchModeDefault, a.GetWebSearchEmulationMode())
}
@@ -0,0 +1,632 @@
//go:build unit
package service
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
)
func TestGrokAccountModelMappingCacheInvalidatesWithRuntimeSettings(t *testing.T) {
original := xai.RuntimeModelMappingOptions()
t.Cleanup(func() { xai.SetRuntimeModelMappingOptions(original) })
account := &Account{Platform: PlatformGrok, Credentials: map[string]any{}}
xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{})
requireMappedModel(t, account, "claude-sonnet-4-5", "claude-sonnet-4-5")
xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{
DefaultText: "grok-build-0.1",
EnableCrossClientMap: true,
})
requireMappedModel(t, account, "claude-sonnet-4-5", "grok-build-0.1")
}
func requireMappedModel(t *testing.T, account *Account, requested, expected string) {
t.Helper()
if actual := account.GetMappedModel(requested); actual != expected {
t.Fatalf("GetMappedModel(%q) = %q, want %q", requested, actual, expected)
}
}
func TestMatchWildcard(t *testing.T) {
tests := []struct {
name string
pattern string
str string
expected bool
}{
// 精确匹配
{"exact match", "claude-sonnet-4-5", "claude-sonnet-4-5", true},
{"exact mismatch", "claude-sonnet-4-5", "claude-opus-4-5", false},
// 通配符匹配
{"wildcard prefix match", "claude-*", "claude-sonnet-4-5", true},
{"wildcard prefix match 2", "claude-*", "claude-opus-4-5-thinking", true},
{"wildcard prefix mismatch", "claude-*", "gemini-3-flash", false},
{"wildcard partial match", "gemini-3*", "gemini-3-flash", true},
{"wildcard partial match 2", "gemini-3*", "gemini-3-pro-image", true},
{"wildcard partial mismatch", "gemini-3*", "gemini-2.5-flash", false},
// 边界情况
{"empty pattern exact", "", "", true},
{"empty pattern mismatch", "", "claude", false},
{"single star", "*", "anything", true},
{"star at end only", "abc*", "abcdef", true},
{"star at end empty suffix", "abc*", "abc", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := matchWildcard(tt.pattern, tt.str)
if result != tt.expected {
t.Errorf("matchWildcard(%q, %q) = %v, want %v", tt.pattern, tt.str, result, tt.expected)
}
})
}
}
func TestMatchWildcardMappingResult(t *testing.T) {
tests := []struct {
name string
mapping map[string]string
requestedModel string
expected string
matched bool
}{
// 精确匹配优先于通配符
{
name: "exact match takes precedence",
mapping: map[string]string{
"claude-sonnet-4-5": "claude-sonnet-4-5-exact",
"claude-*": "claude-default",
},
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-4-5-exact",
matched: true,
},
// 最长通配符优先
{
name: "longer wildcard takes precedence",
mapping: map[string]string{
"claude-*": "claude-default",
"claude-sonnet-*": "claude-sonnet-default",
"claude-sonnet-4*": "claude-sonnet-4-series",
},
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-4-series",
matched: true,
},
// 单个通配符
{
name: "single wildcard",
mapping: map[string]string{
"claude-*": "claude-mapped",
},
requestedModel: "claude-opus-4-5",
expected: "claude-mapped",
matched: true,
},
// 无匹配返回原始模型
{
name: "no match returns original",
mapping: map[string]string{
"claude-*": "claude-mapped",
},
requestedModel: "gemini-3-flash",
expected: "gemini-3-flash",
matched: false,
},
// 空映射返回原始模型
{
name: "empty mapping returns original",
mapping: map[string]string{},
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-4-5",
matched: false,
},
// Gemini 模型映射
{
name: "gemini wildcard mapping",
mapping: map[string]string{
"gemini-3*": "gemini-3-pro-high",
"gemini-2.5*": "gemini-2.5-flash",
},
requestedModel: "gemini-3-flash-preview",
expected: "gemini-3-pro-high",
matched: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, matched := matchWildcardMappingResult(tt.mapping, tt.requestedModel)
if result != tt.expected || matched != tt.matched {
t.Errorf("matchWildcardMappingResult(%v, %q) = (%q, %v), want (%q, %v)", tt.mapping, tt.requestedModel, result, matched, tt.expected, tt.matched)
}
})
}
}
func TestAccountIsModelSupported(t *testing.T) {
tests := []struct {
name string
platform string
credentials map[string]any
requestedModel string
expected bool
}{
// 无映射 = 允许所有
{
name: "no mapping allows all",
credentials: nil,
requestedModel: "any-model",
expected: true,
},
{
name: "empty mapping allows all",
credentials: map[string]any{},
requestedModel: "any-model",
expected: true,
},
// 精确匹配
{
name: "exact match supported",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-sonnet-4-5": "target-model",
},
},
requestedModel: "claude-sonnet-4-5",
expected: true,
},
{
name: "exact match not supported",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-sonnet-4-5": "target-model",
},
},
requestedModel: "claude-opus-4-5",
expected: false,
},
// 通配符匹配
{
name: "wildcard match supported",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-*": "claude-sonnet-4-5",
},
},
requestedModel: "claude-opus-4-5-thinking",
expected: true,
},
{
name: "gemini customtools alias matches normalized mapping",
platform: PlatformGemini,
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3.1-pro-preview": "gemini-3.1-pro-preview",
},
},
requestedModel: "gemini-3.1-pro-preview-customtools",
expected: true,
},
{
name: "wildcard match not supported",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-*": "claude-sonnet-4-5",
},
},
requestedModel: "gemini-3-flash",
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{
Platform: tt.platform,
Credentials: tt.credentials,
}
result := account.IsModelSupported(tt.requestedModel)
if result != tt.expected {
t.Errorf("IsModelSupported(%q) = %v, want %v", tt.requestedModel, result, tt.expected)
}
})
}
}
func TestAccountGetMappedModel(t *testing.T) {
tests := []struct {
name string
platform string
credentials map[string]any
requestedModel string
expected string
}{
// 无映射 = 返回原始模型
{
name: "no mapping returns original",
credentials: nil,
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-4-5",
},
{
name: "no mapping preserves gemini customtools model",
platform: PlatformGemini,
credentials: nil,
requestedModel: "gemini-3.1-pro-preview-customtools",
expected: "gemini-3.1-pro-preview-customtools",
},
// 精确匹配
{
name: "exact match",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-sonnet-4-5": "target-model",
},
},
requestedModel: "claude-sonnet-4-5",
expected: "target-model",
},
// 通配符匹配(最长优先)
{
name: "wildcard longest match",
credentials: map[string]any{
"model_mapping": map[string]any{
"claude-*": "claude-default",
"claude-sonnet-*": "claude-sonnet-mapped",
},
},
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-mapped",
},
// 无匹配返回原始模型
{
name: "gemini customtools alias resolves through normalized mapping",
platform: PlatformGemini,
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3.1-pro-preview": "gemini-3.1-pro-preview",
},
},
requestedModel: "gemini-3.1-pro-preview-customtools",
expected: "gemini-3.1-pro-preview",
},
{
name: "gemini customtools exact mapping wins over normalized fallback",
platform: PlatformGemini,
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3.1-pro-preview": "gemini-3.1-pro-preview",
"gemini-3.1-pro-preview-customtools": "gemini-3.1-pro-preview-customtools",
},
},
requestedModel: "gemini-3.1-pro-preview-customtools",
expected: "gemini-3.1-pro-preview-customtools",
},
{
name: "no match returns original",
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-*": "gemini-mapped",
},
},
requestedModel: "claude-sonnet-4-5",
expected: "claude-sonnet-4-5",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{
Platform: tt.platform,
Credentials: tt.credentials,
}
result := account.GetMappedModel(tt.requestedModel)
if result != tt.expected {
t.Errorf("GetMappedModel(%q) = %q, want %q", tt.requestedModel, result, tt.expected)
}
})
}
}
func TestAccountGetModelMapping_AntigravityNormalizesGemini31ProAliases(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformAntigravity,
Credentials: map[string]any{
"model_mapping": map[string]any{
domain.AntigravityGemini31ProAgentModel: domain.AntigravityGemini31ProAgentModel,
"gemini-3.1-pro-high": "gemini-3.1-pro-high",
"gemini-3.1-pro-preview": "gemini-3.1-pro-high",
},
},
}
mapping := account.GetModelMapping()
if got := mapping["gemini-3.1-pro"]; got != domain.AntigravityGemini31ProAgentModel {
t.Fatalf("expected gemini-3.1-pro to map to %q, got %q", domain.AntigravityGemini31ProAgentModel, got)
}
if got := mapping["gemini-3.1-pro-high"]; got != domain.AntigravityGemini31ProAgentModel {
t.Fatalf("expected gemini-3.1-pro-high to map to %q, got %q", domain.AntigravityGemini31ProAgentModel, got)
}
if got := mapping["gemini-3.1-pro-preview"]; got != domain.AntigravityGemini31ProAgentModel {
t.Fatalf("expected gemini-3.1-pro-preview to map to %q, got %q", domain.AntigravityGemini31ProAgentModel, got)
}
}
func TestAccountGetModelMapping_AntigravityPreservesGemini31ProOverrides(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformAntigravity,
Credentials: map[string]any{
"model_mapping": map[string]any{
domain.AntigravityGemini31ProAgentModel: domain.AntigravityGemini31ProAgentModel,
"gemini-3.1-pro-high": "custom-high",
"gemini-3.1-pro-preview": "custom-preview",
},
},
}
mapping := account.GetModelMapping()
if got := mapping["gemini-3.1-pro-high"]; got != "custom-high" {
t.Fatalf("expected gemini-3.1-pro-high override to be preserved, got %q", got)
}
if got := mapping["gemini-3.1-pro-preview"]; got != "custom-preview" {
t.Fatalf("expected gemini-3.1-pro-preview override to be preserved, got %q", got)
}
if got := mapping["gemini-3.1-pro"]; got != domain.AntigravityGemini31ProAgentModel {
t.Fatalf("expected gemini-3.1-pro alias to default to %q, got %q", domain.AntigravityGemini31ProAgentModel, got)
}
}
func TestAccountGetModelMapping_AntigravityGemini31ProAliasesRespectWildcard(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformAntigravity,
Credentials: map[string]any{
"model_mapping": map[string]any{
domain.AntigravityGemini31ProAgentModel: domain.AntigravityGemini31ProAgentModel,
"gemini-3.1-*": "custom-wildcard",
},
},
}
mapping := account.GetModelMapping()
if got := mapping["gemini-3.1-pro"]; got != "" {
t.Fatalf("expected gemini-3.1-pro exact alias to stay unset when wildcard exists, got %q", got)
}
if got := mapping["gemini-3.1-pro-high"]; got != "" {
t.Fatalf("expected gemini-3.1-pro-high exact alias to stay unset when wildcard exists, got %q", got)
}
if got := mapping["gemini-3.1-pro-preview"]; got != "" {
t.Fatalf("expected gemini-3.1-pro-preview exact alias to stay unset when wildcard exists, got %q", got)
}
}
func TestAccountResolveMappedModel(t *testing.T) {
tests := []struct {
name string
platform string
credentials map[string]any
requestedModel string
expectedModel string
expectedMatch bool
}{
{
name: "no mapping reports unmatched",
credentials: nil,
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: false,
},
{
name: "exact passthrough mapping still counts as matched",
credentials: map[string]any{
"model_mapping": map[string]any{
"gpt-5.4": "gpt-5.4",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: true,
},
{
name: "wildcard passthrough mapping still counts as matched",
credentials: map[string]any{
"model_mapping": map[string]any{
"gpt-*": "gpt-5.4",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: true,
},
{
name: "gemini customtools alias reports normalized match",
platform: PlatformGemini,
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3.1-pro-preview": "gemini-3.1-pro-preview",
},
},
requestedModel: "gemini-3.1-pro-preview-customtools",
expectedModel: "gemini-3.1-pro-preview",
expectedMatch: true,
},
{
name: "gemini customtools exact mapping reports exact match",
platform: PlatformGemini,
credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3.1-pro-preview": "gemini-3.1-pro-preview",
"gemini-3.1-pro-preview-customtools": "gemini-3.1-pro-preview-customtools",
},
},
requestedModel: "gemini-3.1-pro-preview-customtools",
expectedModel: "gemini-3.1-pro-preview-customtools",
expectedMatch: true,
},
{
name: "missing mapping reports unmatched",
credentials: map[string]any{
"model_mapping": map[string]any{
"gpt-5.2": "gpt-5.2",
},
},
requestedModel: "gpt-5.4",
expectedModel: "gpt-5.4",
expectedMatch: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{
Platform: tt.platform,
Credentials: tt.credentials,
}
mappedModel, matched := account.ResolveMappedModel(tt.requestedModel)
if mappedModel != tt.expectedModel || matched != tt.expectedMatch {
t.Fatalf("ResolveMappedModel(%q) = (%q, %v), want (%q, %v)", tt.requestedModel, mappedModel, matched, tt.expectedModel, tt.expectedMatch)
}
})
}
}
func TestAccountGetModelMapping_AntigravityEnsuresGeminiDefaultPassthroughs(t *testing.T) {
account := &Account{
Platform: PlatformAntigravity,
Credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3-pro-high": "gemini-3.1-pro-high",
},
},
}
mapping := account.GetModelMapping()
if mapping["gemini-3-flash"] != "gemini-3-flash" {
t.Fatalf("expected gemini-3-flash passthrough to be auto-filled, got: %q", mapping["gemini-3-flash"])
}
if mapping["gemini-3.1-pro-high"] != "gemini-3.1-pro-high" {
t.Fatalf("expected gemini-3.1-pro-high passthrough to be auto-filled, got: %q", mapping["gemini-3.1-pro-high"])
}
if mapping["gemini-3.1-pro-low"] != "gemini-3.1-pro-low" {
t.Fatalf("expected gemini-3.1-pro-low passthrough to be auto-filled, got: %q", mapping["gemini-3.1-pro-low"])
}
}
func TestAccountGetModelMapping_AntigravityRespectsWildcardOverride(t *testing.T) {
account := &Account{
Platform: PlatformAntigravity,
Credentials: map[string]any{
"model_mapping": map[string]any{
"gemini-3*": "gemini-3.1-pro-high",
},
},
}
mapping := account.GetModelMapping()
if _, exists := mapping["gemini-3-flash"]; exists {
t.Fatalf("did not expect explicit gemini-3-flash passthrough when wildcard already exists")
}
if _, exists := mapping["gemini-3.1-pro-high"]; exists {
t.Fatalf("did not expect explicit gemini-3.1-pro-high passthrough when wildcard already exists")
}
if _, exists := mapping["gemini-3.1-pro-low"]; exists {
t.Fatalf("did not expect explicit gemini-3.1-pro-low passthrough when wildcard already exists")
}
if mapped := account.GetMappedModel("gemini-3-flash"); mapped != "gemini-3.1-pro-high" {
t.Fatalf("expected wildcard mapping to stay effective, got: %q", mapped)
}
}
func TestAccountGetModelMapping_CacheInvalidatesOnCredentialsReplace(t *testing.T) {
account := &Account{
Credentials: map[string]any{
"model_mapping": map[string]any{
"claude-3-5-sonnet": "upstream-a",
},
},
}
first := account.GetModelMapping()
if first["claude-3-5-sonnet"] != "upstream-a" {
t.Fatalf("unexpected first mapping: %v", first)
}
account.Credentials = map[string]any{
"model_mapping": map[string]any{
"claude-3-5-sonnet": "upstream-b",
},
}
second := account.GetModelMapping()
if second["claude-3-5-sonnet"] != "upstream-b" {
t.Fatalf("expected cache invalidated after credentials replace, got: %v", second)
}
}
func TestAccountGetModelMapping_CacheInvalidatesOnMappingLenChange(t *testing.T) {
rawMapping := map[string]any{
"claude-sonnet": "sonnet-a",
}
account := &Account{
Credentials: map[string]any{
"model_mapping": rawMapping,
},
}
first := account.GetModelMapping()
if len(first) != 1 {
t.Fatalf("unexpected first mapping length: %d", len(first))
}
rawMapping["claude-opus"] = "opus-b"
second := account.GetModelMapping()
if second["claude-opus"] != "opus-b" {
t.Fatalf("expected cache invalidated after mapping len change, got: %v", second)
}
}
func TestAccountGetModelMapping_CacheInvalidatesOnInPlaceValueChange(t *testing.T) {
rawMapping := map[string]any{
"claude-sonnet": "sonnet-a",
}
account := &Account{
Credentials: map[string]any{
"model_mapping": rawMapping,
},
}
first := account.GetModelMapping()
if first["claude-sonnet"] != "sonnet-a" {
t.Fatalf("unexpected first mapping: %v", first)
}
rawMapping["claude-sonnet"] = "sonnet-b"
second := account.GetModelMapping()
if second["claude-sonnet"] != "sonnet-b" {
t.Fatalf("expected cache invalidated after in-place value change, got: %v", second)
}
}
@@ -0,0 +1,227 @@
//go:build unit
package service
import (
"bytes"
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func adaptiveProtocolTestAccount(platform string, baseURLs map[string]any) *Account {
return &Account{
ID: 701,
Name: "adaptive-cn",
Platform: platform,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"api_protocol": APIProtocolAdaptive,
"account_mode": AccountModePayG,
"api_base_urls": baseURLs,
},
}
}
func adaptiveProtocolTestContext(path string, body []byte) *gin.Context {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
return c
}
type cnProtocolIngressCase struct {
name string
path string
body []byte
forward func(*OpenAIGatewayService, *gin.Context, *Account, []byte) error
}
func cnProtocolIngressCases() []cnProtocolIngressCase {
return []cnProtocolIngressCase{
{
name: "chat completions",
path: "/v1/chat/completions",
body: []byte(`{"model":"deepseek-chat","messages":[{"role":"user","content":"hello"}],"stream":false}`),
forward: func(svc *OpenAIGatewayService, c *gin.Context, account *Account, body []byte) error {
_, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
return err
},
},
{
name: "messages",
path: "/v1/messages",
body: []byte(`{"model":"deepseek-chat","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`),
forward: func(svc *OpenAIGatewayService, c *gin.Context, account *Account, body []byte) error {
_, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
return err
},
},
{
name: "responses",
path: "/v1/responses",
body: []byte(`{"model":"deepseek-chat","input":"hello","stream":false}`),
forward: func(svc *OpenAIGatewayService, c *gin.Context, account *Account, body []byte) error {
_, err := svc.Forward(context.Background(), c, account, body)
return err
},
},
}
}
func TestAdaptiveProtocolRoutesChatCompletionsToNativeChat(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"glm-4.7","messages":[{"role":"user","content":"hello"}],"stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformZhipu, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
})
_, err := svc.ForwardAsChatCompletions(context.Background(), adaptiveProtocolTestContext("/v1/chat/completions", body), account, body, "", "")
require.Error(t, err)
require.Equal(t, "http://chat.example/v1/chat/completions", upstream.lastReq.URL.String())
require.True(t, gjson.GetBytes(upstream.lastBody, "messages").IsArray())
require.False(t, gjson.GetBytes(upstream.lastBody, "input").Exists())
}
func TestAdaptiveProtocolRoutesResponsesShapedChatToNativeResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"deepseek-v4","input":"hello","max_output_tokens":32,"stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformDeepseek, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
APIProtocolResponses: "http://responses.example",
})
_, err := svc.ForwardAsChatCompletions(context.Background(), adaptiveProtocolTestContext("/v1/chat/completions", body), account, body, "", "")
require.Error(t, err)
require.Equal(t, "http://responses.example/responses", upstream.lastReq.URL.String())
require.True(t, gjson.GetBytes(upstream.lastBody, "input").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "messages").Exists())
}
func TestAdaptiveProtocolConvertsResponsesShapedChatForChatOnlyProvider(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"kimi-k2.5","input":"hello","max_output_tokens":32,"stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformKimi, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
})
_, err := svc.ForwardAsChatCompletions(context.Background(), adaptiveProtocolTestContext("/v1/chat/completions", body), account, body, "", "")
require.Error(t, err)
require.Equal(t, "http://chat.example/v1/chat/completions", upstream.lastReq.URL.String())
require.True(t, gjson.GetBytes(upstream.lastBody, "messages").IsArray())
require.False(t, gjson.GetBytes(upstream.lastBody, "input").Exists())
}
func TestAdaptiveProtocolRoutesMessagesToNativeAnthropic(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"glm-4.7","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformZhipu, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
})
_, err := svc.ForwardAsAnthropic(context.Background(), adaptiveProtocolTestContext("/v1/messages", body), account, body, "", "")
require.Error(t, err)
require.Equal(t, "http://anthropic.example/v1/messages", upstream.lastReq.URL.String())
require.Equal(t, "glm-4.7", gjson.GetBytes(upstream.lastBody, "model").String())
}
func TestAdaptiveProtocolConvertsKimiResponsesToChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"kimi-k2.5","input":"hello","stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformKimi, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
})
_, err := svc.Forward(context.Background(), adaptiveProtocolTestContext("/v1/responses", body), account, body)
require.Error(t, err)
require.Equal(t, "http://chat.example/v1/chat/completions", upstream.lastReq.URL.String())
require.True(t, gjson.GetBytes(upstream.lastBody, "messages").IsArray())
require.False(t, gjson.GetBytes(upstream.lastBody, "input").Exists())
}
func TestAdaptiveProtocolRoutesDeepSeekResponsesToNativeResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"deepseek-v4","input":"hello","max_output_tokens":32,"store":true,"previous_response_id":"resp_old","stream":false}`)
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformDeepseek, map[string]any{
APIProtocolChatCompletions: "http://chat.example",
APIProtocolAnthropic: "http://anthropic.example",
APIProtocolResponses: "http://responses.example",
})
_, err := svc.Forward(context.Background(), adaptiveProtocolTestContext("/v1/responses", body), account, body)
require.Error(t, err)
require.Equal(t, "http://responses.example/responses", upstream.lastReq.URL.String())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "previous_response_id").Exists())
require.Equal(t, int64(32), gjson.GetBytes(upstream.lastBody, "max_output_tokens").Int())
require.False(t, gjson.GetBytes(upstream.lastBody, "instructions").Exists())
}
func TestFixedCNChatProtocolOverridesStaleResponsesMode(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range cnProtocolIngressCases() {
t.Run(tc.name, func(t *testing.T) {
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformDeepseek, nil)
account.Credentials["api_protocol"] = APIProtocolChatCompletions
account.Credentials["base_url"] = "http://chat.example"
account.Extra = map[string]any{
openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeForceResponses),
}
err := tc.forward(svc, adaptiveProtocolTestContext(tc.path, tc.body), account, tc.body)
require.Error(t, err)
require.Equal(t, "http://chat.example/v1/chat/completions", upstream.lastReq.URL.String())
})
}
}
func TestFixedCNResponsesProtocolOverridesStaleChatMode(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range cnProtocolIngressCases() {
t.Run(tc.name, func(t *testing.T) {
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := adaptiveProtocolTestAccount(PlatformDeepseek, nil)
account.Credentials["api_protocol"] = APIProtocolResponses
account.Credentials["base_url"] = "http://responses.example"
account.Extra = map[string]any{
openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeForceChatCompletions),
}
err := tc.forward(svc, adaptiveProtocolTestContext(tc.path, tc.body), account, tc.body)
require.Error(t, err)
require.Equal(t, "http://responses.example/responses", upstream.lastReq.URL.String())
})
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,20 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestNormalizeAccountConcurrencyDefaultsInvalidGrokOAuthToOne(t *testing.T) {
require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 0))
require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, -5))
}
func TestNormalizeAccountConcurrencyPreservesExplicitValues(t *testing.T) {
require.Equal(t, 50, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 50))
require.Equal(t, 2, normalizeAccountConcurrency(PlatformOpenAI, AccountTypeOAuth, 2))
require.Equal(t, 2, normalizeAccountConcurrency(PlatformGrok, AccountTypeAPIKey, 2))
}
@@ -0,0 +1,751 @@
package service
import (
"context"
"net/http"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
type upstreamBillingProbeAdminRepo struct {
*upstreamBillingProbeAccountRepo
}
func (r *upstreamBillingProbeAdminRepo) ListShadowsByParent(context.Context, int64) ([]*Account, error) {
return nil, nil
}
type accountBillingSettingsAdminRepo struct {
*upstreamBillingProbeAccountRepo
concurrentRate *float64
lastExplicitRate *float64
updateCalls int
}
func (r *accountBillingSettingsAdminRepo) UpdateWithAccountBillingSettings(
_ context.Context,
account *Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error {
r.mu.Lock()
defer r.mu.Unlock()
current := r.accounts[account.ID]
if current == nil {
return ErrAccountNotFound
}
updated := *account
updated.Credentials = mergeMap(nil, account.Credentials)
updated.Extra = mergeMap(nil, account.Extra)
if updated.Extra == nil {
updated.Extra = make(map[string]any)
}
if probeEnabled != nil {
updated.Extra[UpstreamBillingProbeEnabledExtraKey] = *probeEnabled
}
if rateSyncEnabled != nil {
updated.Extra[UpstreamBillingRateSyncEnabledExtraKey] = *rateSyncEnabled
}
switch {
case rateMultiplier != nil:
value := *rateMultiplier
updated.RateMultiplier = &value
r.lastExplicitRate = &value
case r.concurrentRate != nil:
value := *r.concurrentRate
updated.RateMultiplier = &value
r.lastExplicitRate = nil
default:
updated.RateMultiplier = cloneAccountValuePointer(current.RateMultiplier)
r.lastExplicitRate = nil
}
r.accounts[account.ID] = &updated
r.updateCalls++
return nil
}
func TestUpdateAccountRoutesRateIntentThroughAtomicBillingUpdater(t *testing.T) {
accountID := int64(109)
initialRate := 0.1
concurrentRate := 0.2
repo := &accountBillingSettingsAdminRepo{
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Name: "before",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
RateMultiplier: &initialRate,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
},
}},
concurrentRate: &concurrentRate,
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{Name: "after"})
require.NoError(t, err)
require.Equal(t, 1, repo.updateCalls)
require.Nil(t, repo.lastExplicitRate)
require.Equal(t, concurrentRate, *updated.RateMultiplier)
// 手工倍率只有在同步不再开启时才被接受,所以同一请求先关闭同步再设值
// (同步仍开启时的手工倍率由 TestUpdateAccountRejectsManualRateWhileRateSyncEnabled 覆盖)。
zero := 0.0
syncDisabled := false
updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &syncDisabled,
RateMultiplier: &zero,
})
require.NoError(t, err)
require.Equal(t, 2, repo.updateCalls)
require.NotNil(t, repo.lastExplicitRate)
require.Zero(t, *repo.lastExplicitRate)
require.Zero(t, *updated.RateMultiplier)
}
func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
svc := &adminServiceImpl{accountRepo: repo}
created, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
Name: "upstream",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "sk-test"},
SkipDefaultGroupBind: true,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
})
require.NoError(t, err)
require.NotContains(t, created.Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, created.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey)
}
func TestCreateAccountAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) {
enabled := true
repo := &upstreamBillingProbeAccountRepo{}
created, err := (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{
Name: "upstream",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "sk-test"},
ProbeEnabled: &enabled,
SkipDefaultGroupBind: true,
})
require.NoError(t, err)
require.Equal(t, true, created.Extra[UpstreamBillingProbeEnabledExtraKey])
_, err = (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{
Name: "oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{"access_token": "token"},
ProbeEnabled: &enabled,
SkipDefaultGroupBind: true,
})
require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid)
}
func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(t *testing.T) {
accountID := int64(110)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{"custom": "value"},
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
require.Contains(t, updated.Extra, UpstreamBillingProbeExtraKey)
require.Equal(t, "value", updated.Extra["custom"])
}
func TestUpdateAccountPreservesGrokBillingSnapshotForUnrelatedEdit(t *testing.T) {
accountID := int64(112)
billing := &xai.BillingSummary{
StatusCode: http.StatusForbidden,
WeeklyStatusCode: http.StatusForbidden,
}
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{grokBillingExtraKey: billing},
},
}}
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{"custom": "value"},
})
require.NoError(t, err)
require.Equal(t, billing, updated.Extra[grokBillingExtraKey])
require.Equal(t, "value", updated.Extra["custom"])
eligible, reason := updated.GrokMediaGenerationEligibility()
require.False(t, eligible)
require.Equal(t, "billing_forbidden", reason)
}
func TestUpdateAccountPreservesProbeSnapshotWhenIdentityValuesAreUnchanged(t *testing.T) {
accountID := int64(119)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Credentials: map[string]any{
"api_key": "sk-existing",
"base_url": "https://upstream.example",
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"x-route": "stable"},
},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Credentials: map[string]any{
"base_url": "https://upstream.example",
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"x-route": "stable"},
},
})
require.NoError(t, err)
require.Contains(t, updated.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpdateAccountInvalidatesProbeSnapshotWhenUpstreamIdentityChanges(t *testing.T) {
tests := []struct {
name string
input *UpdateAccountInput
wantEnabled bool
}{
{
name: "api key",
input: &UpdateAccountInput{Credentials: map[string]any{"api_key": "sk-new"}},
wantEnabled: true,
},
{
name: "base url",
input: &UpdateAccountInput{Credentials: map[string]any{"base_url": "https://new.example"}},
wantEnabled: true,
},
{
name: "header override",
input: &UpdateAccountInput{Credentials: map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"x-route": "new"},
}},
wantEnabled: true,
},
{
name: "account type",
input: &UpdateAccountInput{Type: AccountTypeOAuth},
wantEnabled: false,
},
}
for i, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
accountID := int64(120 + i)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Credentials: map[string]any{
"api_key": "sk-old",
"base_url": "https://old.example",
},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, tt.input)
require.NoError(t, err)
require.NotContains(t, updated.Extra, UpstreamBillingProbeExtraKey)
if tt.wantEnabled {
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
} else {
require.NotContains(t, updated.Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey)
}
})
}
}
func TestUpdateAccountInvalidatesProbeSnapshotWhenProxyChanges(t *testing.T) {
accountID := int64(140)
oldProxyID := int64(7)
newProxyID := int64(8)
baseRepo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Credentials: map[string]any{"api_key": "sk-test"},
ProxyID: &oldProxyID,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
updated, err := (&adminServiceImpl{accountRepo: &upstreamBillingProbeAdminRepo{baseRepo}}).UpdateAccount(
context.Background(),
accountID,
&UpdateAccountInput{ProxyID: &newProxyID},
)
require.NoError(t, err)
require.Equal(t, newProxyID, *updated.ProxyID)
require.NotContains(t, updated.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpdateAccountPreservesProbeSnapshotWhenProxyIsUnchanged(t *testing.T) {
accountID := int64(141)
existingProxyID := int64(7)
unchangedProxyID := int64(7)
baseRepo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Credentials: map[string]any{"api_key": "sk-test"},
ProxyID: &existingProxyID,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
updated, err := (&adminServiceImpl{accountRepo: &upstreamBillingProbeAdminRepo{baseRepo}}).UpdateAccount(
context.Background(),
accountID,
&UpdateAccountInput{ProxyID: &unchangedProxyID},
)
require.NoError(t, err)
require.Contains(t, updated.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpdateAccountAcceptsProbeEnabledAndRejectsInjectedSnapshot(t *testing.T) {
accountID := int64(111)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, updated.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpdateAccountRateSyncControlsProbeAndManualMode(t *testing.T) {
accountID := int64(151)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
syncEnabled := true
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &syncEnabled,
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
syncEnabled = false
updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &syncEnabled,
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
}
// 单账号编辑必须和批量路径语义一致:同步开启时倍率归上游所有,手工值会在下一次
// 成功探测时被覆盖,因此直接拒绝而不是静默接受。
func TestUpdateAccountRejectsManualRateWhileRateSyncEnabled(t *testing.T) {
newRepo := func(accountID int64, extra map[string]any) *upstreamBillingProbeAccountRepo {
initialRate := 0.25
return &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
RateMultiplier: &initialRate,
Extra: extra,
},
}}
}
manualRate := 3.5
syncEnabled := map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
}
t.Run("sync enabled rejects manual rate", func(t *testing.T) {
accountID := int64(153)
repo := newRepo(accountID, mergeMap(nil, syncEnabled))
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateMultiplier: &manualRate,
})
require.ErrorIs(t, err, ErrUpstreamBillingRateSyncConflict)
require.Equal(t, 0.25, *repo.accounts[accountID].RateMultiplier)
})
t.Run("enabling sync in the same request rejects manual rate", func(t *testing.T) {
accountID := int64(154)
repo := newRepo(accountID, map[string]any{})
enable := true
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &enable,
RateMultiplier: &manualRate,
})
require.ErrorIs(t, err, ErrUpstreamBillingRateSyncConflict)
require.Equal(t, 0.25, *repo.accounts[accountID].RateMultiplier)
})
// 用户显式收回所有权:同一请求关闭同步并改倍率必须放行。
t.Run("disabling sync in the same request allows manual rate", func(t *testing.T) {
accountID := int64(155)
repo := newRepo(accountID, mergeMap(nil, syncEnabled))
disable := false
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &disable,
RateMultiplier: &manualRate,
})
require.NoError(t, err)
require.Equal(t, false, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
require.NotNil(t, updated.RateMultiplier)
require.Equal(t, manualRate, *updated.RateMultiplier)
})
t.Run("sync disabled allows manual rate", func(t *testing.T) {
accountID := int64(156)
repo := newRepo(accountID, map[string]any{UpstreamBillingProbeEnabledExtraKey: true})
updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateMultiplier: &manualRate,
})
require.NoError(t, err)
require.NotNil(t, updated.RateMultiplier)
require.Equal(t, manualRate, *updated.RateMultiplier)
})
}
func TestUpdateAccountRejectsSyncWithExplicitlyDisabledProbe(t *testing.T) {
accountID := int64(152)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Status: StatusActive,
},
}}
probeEnabled := false
syncEnabled := true
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
ProbeEnabled: &probeEnabled,
RateSyncEnabled: &syncEnabled,
})
require.Error(t, err)
require.Empty(t, repo.updates[accountID])
}
func TestUpdateAccountExplicitProbeDisableUsesDedicatedExtraUpdate(t *testing.T) {
accountID := int64(113)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: false},
})
require.NoError(t, err)
require.Len(t, repo.updates[accountID], 1)
require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingRateSyncEnabledExtraKey])
}
func TestUpdateAccountExplicitUnchangedProbeEnabledStillUsesDedicatedExtraUpdate(t *testing.T) {
accountID := int64(114)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: true},
},
}}
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: true},
})
require.NoError(t, err)
require.Len(t, repo.updates[accountID], 1)
require.Equal(t, true, repo.updates[accountID][0][UpstreamBillingProbeEnabledExtraKey])
}
func TestUpdateAccountRejectsInvalidProbeEnabled(t *testing.T) {
accountID := int64(112)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
_, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: "true"},
})
require.Error(t, err)
}
func TestUpdateAccountExtraDropsManagedBillingProbeFields(t *testing.T) {
accountID := int64(153)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
err := (&adminServiceImpl{accountRepo: repo}).UpdateAccountExtra(context.Background(), accountID, map[string]any{
"custom": "value",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
})
require.NoError(t, err)
require.Equal(t, "value", repo.accounts[accountID].Extra["custom"])
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeExtraKey)
}
func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
svc := &adminServiceImpl{accountRepo: repo}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{
"custom": "value",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
}
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Equal(t, 1, result.Success)
require.Len(t, repo.bulkUpdates, 1)
require.Equal(t, "value", repo.bulkUpdates[0].Extra["custom"])
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey)
}
func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) {
for _, enabled := range []bool{true, false} {
t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
ProbeEnabled: &enabled,
})
require.NoError(t, err)
require.Equal(t, 2, result.Success)
require.Len(t, repo.bulkUpdates, 1)
require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey])
if !enabled {
require.Equal(t, false, repo.bulkUpdates[0].Extra[UpstreamBillingRateSyncEnabledExtraKey])
}
require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled)
require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled)
})
}
}
func TestBulkUpdateAccountsRejectsProbeSettingForIneligibleTargetBeforeWrite(t *testing.T) {
for _, enabled := range []bool{true, false} {
t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
}}
_, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
ProbeEnabled: &enabled,
})
require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid)
require.Empty(t, repo.bulkUpdates)
})
}
}
func TestBulkUpdateAccountsRejectsProbeSettingWhenTargetIsMissing(t *testing.T) {
enabled := true
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
_, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
ProbeEnabled: &enabled,
})
require.ErrorIs(t, err, ErrAccountNotFound)
require.Empty(t, repo.bulkUpdates)
}
func TestBulkUpdateAccountsInvalidatesProbeSnapshotForIdentityCredentials(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: map[string]any{"api_key": "sk-new"},
}
result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Equal(t, 1, result.Success)
require.Len(t, repo.bulkUpdates, 1)
require.Contains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey)
require.Nil(t, repo.bulkUpdates[0].Extra[UpstreamBillingProbeExtraKey])
}
func TestBulkUpdateAccountsInvalidatesProbeSnapshotForProxyUpdate(t *testing.T) {
proxyID := int64(9)
baseRepo := &upstreamBillingProbeAccountRepo{}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
ProxyID: &proxyID,
}
result, err := (&adminServiceImpl{accountRepo: &upstreamBillingProbeAdminRepo{baseRepo}}).BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Equal(t, 1, result.Success)
require.Len(t, baseRepo.bulkUpdates, 1)
require.Contains(t, baseRepo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey)
require.Nil(t, baseRepo.bulkUpdates[0].Extra[UpstreamBillingProbeExtraKey])
}
func TestBulkUpdateAccountsKeepsProbeSnapshotForUnrelatedCredentials(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: map[string]any{"model_mapping": map[string]any{"gpt-old": "gpt-new"}},
}
_, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Len(t, repo.bulkUpdates, 1)
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey)
}
@@ -0,0 +1,86 @@
package service
import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
func TestMergeBalanceHistoryCodesIncludesAffiliateTransfersByDefault(t *testing.T) {
t.Parallel()
now := time.Date(2026, 5, 3, 12, 0, 0, 0, time.UTC)
older := now.Add(-2 * time.Hour)
newer := now.Add(time.Hour)
usedBy := int64(10)
redeemCodes := []RedeemCode{
{
ID: 1,
Type: RedeemTypeBalance,
Value: 8,
Status: StatusUsed,
UsedBy: &usedBy,
UsedAt: &now,
CreatedAt: now,
},
{
ID: 2,
Type: RedeemTypeConcurrency,
Value: 1,
Status: StatusUsed,
UsedBy: &usedBy,
UsedAt: &older,
CreatedAt: older,
},
}
affiliateCodes := []RedeemCode{
{
ID: -20,
Type: RedeemTypeAffiliateBalance,
Value: 3.5,
Status: StatusUsed,
UsedBy: &usedBy,
UsedAt: &newer,
CreatedAt: newer,
},
}
got := mergeBalanceHistoryCodes(redeemCodes, affiliateCodes, pagination.PaginationParams{
Page: 1,
PageSize: 2,
})
require.Len(t, got, 2)
require.Equal(t, RedeemTypeAffiliateBalance, got[0].Type)
require.Equal(t, RedeemTypeBalance, got[1].Type)
}
func TestMergeBalanceHistoryCodesPaginatesAfterCombiningSources(t *testing.T) {
t.Parallel()
base := time.Date(2026, 5, 3, 12, 0, 0, 0, time.UTC)
usedBy := int64(10)
at := func(hours int) *time.Time {
v := base.Add(time.Duration(hours) * time.Hour)
return &v
}
got := mergeBalanceHistoryCodes(
[]RedeemCode{
{ID: 1, Type: RedeemTypeBalance, UsedBy: &usedBy, UsedAt: at(4), CreatedAt: *at(4)},
{ID: 2, Type: RedeemTypeConcurrency, UsedBy: &usedBy, UsedAt: at(2), CreatedAt: *at(2)},
},
[]RedeemCode{
{ID: -3, Type: RedeemTypeAffiliateBalance, UsedBy: &usedBy, UsedAt: at(3), CreatedAt: *at(3)},
{ID: -4, Type: RedeemTypeAffiliateBalance, UsedBy: &usedBy, UsedAt: at(1), CreatedAt: *at(1)},
},
pagination.PaginationParams{Page: 2, PageSize: 2},
)
require.Len(t, got, 2)
require.Equal(t, RedeemTypeConcurrency, got[0].Type)
require.Equal(t, int64(-4), got[1].ID)
}
@@ -0,0 +1,161 @@
package service
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strconv"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
const (
AdminComplianceVersion = "v2026.06.10"
AdminComplianceDocumentPathZH = "docs/legal/admin-compliance.zh.md"
AdminComplianceDocumentPathEN = "docs/legal/admin-compliance.en.md"
AdminComplianceDocumentURLZH = "https://github.com/Wei-Shaw/sub2api/blob/main/docs/legal/admin-compliance.zh.md"
AdminComplianceDocumentURLEN = "https://github.com/Wei-Shaw/sub2api/blob/main/docs/legal/admin-compliance.en.md"
AdminComplianceAckPhraseZH = "我已阅读、理解并同意 Sub2API 部署与运营合规承诺"
AdminComplianceAckPhraseEN = "I have read, understood, and agree to the Sub2API Deployment and Operation Compliance Commitment"
settingKeyAdminComplianceAcknowledgement = "admin_compliance_acknowledgement"
)
var (
ErrAdminComplianceAcknowledgementRequired = infraerrors.New(
http.StatusLocked,
"ADMIN_COMPLIANCE_ACK_REQUIRED",
"administrator compliance acknowledgement is required",
)
ErrAdminComplianceInvalidPhrase = infraerrors.BadRequest(
"ADMIN_COMPLIANCE_INVALID_PHRASE",
"confirmation phrase does not match",
)
)
type AdminComplianceAcknowledgement struct {
Version string `json:"version"`
DocumentZH string `json:"document_zh"`
DocumentEN string `json:"document_en"`
AdminUserID int64 `json:"admin_user_id"`
IPAddress string `json:"ip_address,omitempty"`
UserAgent string `json:"user_agent,omitempty"`
AcceptedAt time.Time `json:"accepted_at"`
}
type AdminComplianceStatus struct {
Required bool `json:"required"`
Version string `json:"version"`
DocumentPathZH string `json:"document_path_zh"`
DocumentPathEN string `json:"document_path_en"`
DocumentURLZH string `json:"document_url_zh"`
DocumentURLEN string `json:"document_url_en"`
AckPhraseZH string `json:"ack_phrase_zh"`
AckPhraseEN string `json:"ack_phrase_en"`
Acknowledgement *AdminComplianceAcknowledgement `json:"acknowledgement,omitempty"`
}
type AdminComplianceAcceptInput struct {
AdminUserID int64
Phrase string
Language string
IPAddress string
UserAgent string
}
func normalizeAdminComplianceLanguage(raw string) string {
raw = strings.ToLower(strings.TrimSpace(raw))
if strings.HasPrefix(raw, "zh") {
return "zh"
}
return "en"
}
func expectedAdminCompliancePhrase(language string) string {
if normalizeAdminComplianceLanguage(language) == "zh" {
return AdminComplianceAckPhraseZH
}
return AdminComplianceAckPhraseEN
}
func adminComplianceAcknowledgementKey(adminUserID int64) string {
if adminUserID <= 0 {
return settingKeyAdminComplianceAcknowledgement
}
return settingKeyAdminComplianceAcknowledgement + ":" + strconv.FormatInt(adminUserID, 10)
}
func (s *SettingService) GetAdminComplianceStatus(ctx context.Context, adminUserID int64) (*AdminComplianceStatus, error) {
status := &AdminComplianceStatus{
Required: true,
Version: AdminComplianceVersion,
DocumentPathZH: AdminComplianceDocumentPathZH,
DocumentPathEN: AdminComplianceDocumentPathEN,
DocumentURLZH: AdminComplianceDocumentURLZH,
DocumentURLEN: AdminComplianceDocumentURLEN,
AckPhraseZH: AdminComplianceAckPhraseZH,
AckPhraseEN: AdminComplianceAckPhraseEN,
}
if s == nil || s.settingRepo == nil {
return status, nil
}
raw, err := s.settingRepo.GetValue(ctx, adminComplianceAcknowledgementKey(adminUserID))
if err != nil {
if errors.Is(err, ErrSettingNotFound) {
return status, nil
}
return nil, fmt.Errorf("get admin compliance acknowledgement: %w", err)
}
var ack AdminComplianceAcknowledgement
if err := json.Unmarshal([]byte(raw), &ack); err != nil {
return status, nil
}
if ack.Version == AdminComplianceVersion {
status.Required = false
status.Acknowledgement = &ack
}
return status, nil
}
func (s *SettingService) IsAdminComplianceAcknowledged(ctx context.Context, adminUserID int64) (bool, error) {
status, err := s.GetAdminComplianceStatus(ctx, adminUserID)
if err != nil {
return false, err
}
return status != nil && !status.Required, nil
}
func (s *SettingService) AcceptAdminCompliance(ctx context.Context, input AdminComplianceAcceptInput) (*AdminComplianceStatus, error) {
if s == nil || s.settingRepo == nil {
return nil, infraerrors.InternalServer("SETTING_SERVICE_UNAVAILABLE", "setting service is unavailable")
}
phrase := strings.TrimSpace(input.Phrase)
if phrase != expectedAdminCompliancePhrase(input.Language) {
return nil, ErrAdminComplianceInvalidPhrase
}
ack := AdminComplianceAcknowledgement{
Version: AdminComplianceVersion,
DocumentZH: AdminComplianceDocumentPathZH,
DocumentEN: AdminComplianceDocumentPathEN,
AdminUserID: input.AdminUserID,
IPAddress: strings.TrimSpace(input.IPAddress),
UserAgent: strings.TrimSpace(input.UserAgent),
AcceptedAt: time.Now().UTC(),
}
payload, err := json.Marshal(ack)
if err != nil {
return nil, fmt.Errorf("marshal admin compliance acknowledgement: %w", err)
}
if err := s.settingRepo.Set(ctx, adminComplianceAcknowledgementKey(input.AdminUserID), string(payload)); err != nil {
return nil, fmt.Errorf("save admin compliance acknowledgement: %w", err)
}
return s.GetAdminComplianceStatus(ctx, input.AdminUserID)
}
@@ -0,0 +1,133 @@
package service
import (
"context"
"encoding/json"
"errors"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type adminComplianceRepoStub struct {
values map[string]string
}
func (r *adminComplianceRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
if value, ok := r.values[key]; ok {
return &Setting{Key: key, Value: value}, nil
}
return nil, ErrSettingNotFound
}
func (r *adminComplianceRepoStub) GetValue(ctx context.Context, key string) (string, error) {
setting, err := r.Get(ctx, key)
if err != nil {
return "", err
}
return setting.Value, nil
}
func (r *adminComplianceRepoStub) Set(ctx context.Context, key, value string) error {
if r.values == nil {
r.values = map[string]string{}
}
r.values[key] = value
return nil
}
func (r *adminComplianceRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *adminComplianceRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
return nil
}
func (r *adminComplianceRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *adminComplianceRepoStub) Delete(ctx context.Context, key string) error {
delete(r.values, key)
return nil
}
func TestAdminComplianceStatusRequiresAckWhenMissing(t *testing.T) {
svc := NewSettingService(&adminComplianceRepoStub{}, &config.Config{})
status, err := svc.GetAdminComplianceStatus(context.Background(), 1)
require.NoError(t, err)
require.True(t, status.Required)
require.Equal(t, AdminComplianceVersion, status.Version)
require.Equal(t, AdminComplianceAckPhraseZH, status.AckPhraseZH)
require.Equal(t, AdminComplianceDocumentPathZH, status.DocumentPathZH)
}
func TestAcceptAdminComplianceRejectsWrongPhrase(t *testing.T) {
svc := NewSettingService(&adminComplianceRepoStub{}, &config.Config{})
_, err := svc.AcceptAdminCompliance(context.Background(), AdminComplianceAcceptInput{
AdminUserID: 1,
Language: "zh",
Phrase: "我同意",
})
require.Error(t, err)
require.True(t, errors.Is(err, ErrAdminComplianceInvalidPhrase))
}
func TestAcceptAdminCompliancePersistsCurrentVersion(t *testing.T) {
repo := &adminComplianceRepoStub{}
svc := NewSettingService(repo, &config.Config{})
status, err := svc.AcceptAdminCompliance(context.Background(), AdminComplianceAcceptInput{
AdminUserID: 42,
Language: "zh-CN",
Phrase: AdminComplianceAckPhraseZH,
IPAddress: "203.0.113.10",
UserAgent: "test-agent",
})
require.NoError(t, err)
require.False(t, status.Required)
require.NotNil(t, status.Acknowledgement)
require.Equal(t, int64(42), status.Acknowledgement.AdminUserID)
require.Equal(t, "203.0.113.10", status.Acknowledgement.IPAddress)
var stored AdminComplianceAcknowledgement
require.NoError(t, json.Unmarshal([]byte(repo.values[adminComplianceAcknowledgementKey(42)]), &stored))
require.Equal(t, AdminComplianceVersion, stored.Version)
require.Equal(t, AdminComplianceDocumentPathZH, stored.DocumentZH)
}
func TestAdminComplianceStatusRequiresAckOnOldVersion(t *testing.T) {
old, err := json.Marshal(AdminComplianceAcknowledgement{Version: "v2026.01.01"})
require.NoError(t, err)
svc := NewSettingService(&adminComplianceRepoStub{
values: map[string]string{adminComplianceAcknowledgementKey(1): string(old)},
}, &config.Config{})
status, err := svc.GetAdminComplianceStatus(context.Background(), 1)
require.NoError(t, err)
require.True(t, status.Required)
require.Nil(t, status.Acknowledgement)
}
func TestAdminComplianceStatusIsPerAdminUser(t *testing.T) {
current, err := json.Marshal(AdminComplianceAcknowledgement{
Version: AdminComplianceVersion,
AdminUserID: 1,
})
require.NoError(t, err)
svc := NewSettingService(&adminComplianceRepoStub{
values: map[string]string{adminComplianceAcknowledgementKey(1): string(current)},
}, &config.Config{})
statusForUserOne, err := svc.GetAdminComplianceStatus(context.Background(), 1)
require.NoError(t, err)
require.False(t, statusForUserOne.Required)
statusForUserTwo, err := svc.GetAdminComplianceStatus(context.Background(), 2)
require.NoError(t, err)
require.True(t, statusForUserTwo.Required)
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,232 @@
package service
import (
"context"
"crypto/sha256"
"errors"
"fmt"
"strconv"
"strings"
"time"
)
const (
maxGroupNameRunes = 100
duplicateGroupInactiveStatus = "inactive"
)
func duplicateGroupOperationID(sourceID int64, actorScope, operationKey string) string {
operationKey = strings.TrimSpace(operationKey)
if operationKey == "" {
return ""
}
actorScope = strings.TrimSpace(actorScope)
if actorScope == "" {
actorScope = "admin:0"
}
payload := "admin.groups.duplicate\x00" + actorScope + "\x00" + strconv.FormatInt(sourceID, 10) + "\x00" + operationKey
digest := sha256.Sum256([]byte(payload))
return fmt.Sprintf("%x", digest)
}
func duplicateGroupName(sourceName string, copyNumber int) string {
if copyNumber < 1 {
copyNumber = 1
}
suffix := " (Copy)"
if copyNumber > 1 {
suffix = fmt.Sprintf(" (Copy %d)", copyNumber)
}
baseRunes := []rune(strings.TrimSpace(sourceName))
maxBaseRunes := maxGroupNameRunes - len([]rune(suffix))
if maxBaseRunes < 0 {
maxBaseRunes = 0
}
if len(baseRunes) > maxBaseRunes {
baseRunes = baseRunes[:maxBaseRunes]
}
return string(baseRunes) + suffix
}
func cloneGroupValuePointer[T any](value *T) *T {
if value == nil {
return nil
}
cloned := *value
return &cloned
}
func cloneGroupModelRouting(value map[string][]int64) map[string][]int64 {
if value == nil {
return nil
}
cloned := make(map[string][]int64, len(value))
for model, accountIDs := range value {
cloned[model] = append([]int64(nil), accountIDs...)
}
return cloned
}
func cloneGroupVideoModelPrices(value map[string]map[string]float64) map[string]map[string]float64 {
if value == nil {
return nil
}
cloned := make(map[string]map[string]float64, len(value))
for model, prices := range value {
clonedPrices := make(map[string]float64, len(prices))
for resolution, price := range prices {
clonedPrices[resolution] = price
}
cloned[model] = clonedPrices
}
return cloned
}
func cloneGroupMessagesDispatchModelConfig(value OpenAIMessagesDispatchModelConfig) OpenAIMessagesDispatchModelConfig {
cloned := value
if value.ExactModelMappings != nil {
cloned.ExactModelMappings = make(map[string]string, len(value.ExactModelMappings))
for requestedModel, mappedModel := range value.ExactModelMappings {
cloned.ExactModelMappings[requestedModel] = mappedModel
}
}
return cloned
}
func cloneGroupForDuplicate(source *Group, operationID string) *Group {
return &Group{
Name: duplicateGroupName(source.Name, 1),
Description: source.Description,
Platform: source.Platform,
RateMultiplier: source.RateMultiplier,
PeakRateEnabled: source.PeakRateEnabled,
PeakStart: source.PeakStart,
PeakEnd: source.PeakEnd,
PeakRateMultiplier: source.PeakRateMultiplier,
ProfitControlEnabled: source.ProfitControlEnabled,
ProfitMinMargin: source.ProfitMinMargin,
ProfitSafetyBuffer: source.ProfitSafetyBuffer,
IsExclusive: source.IsExclusive,
Status: duplicateGroupInactiveStatus,
DuplicateOperationID: operationID,
SubscriptionType: source.SubscriptionType,
DailyLimitUSD: cloneGroupValuePointer(source.DailyLimitUSD),
WeeklyLimitUSD: cloneGroupValuePointer(source.WeeklyLimitUSD),
MonthlyLimitUSD: cloneGroupValuePointer(source.MonthlyLimitUSD),
DefaultValidityDays: source.DefaultValidityDays,
AllowImageGeneration: source.AllowImageGeneration,
AllowBatchImageGeneration: source.AllowBatchImageGeneration,
ImageRateIndependent: source.ImageRateIndependent,
ImageRateMultiplier: source.ImageRateMultiplier,
ImagePrice1K: cloneGroupValuePointer(source.ImagePrice1K),
ImagePrice2K: cloneGroupValuePointer(source.ImagePrice2K),
ImagePrice4K: cloneGroupValuePointer(source.ImagePrice4K),
BatchImageDiscountMultiplier: source.BatchImageDiscountMultiplier,
BatchImageHoldMultiplier: source.BatchImageHoldMultiplier,
VideoRateIndependent: source.VideoRateIndependent,
VideoRateMultiplier: source.VideoRateMultiplier,
VideoPrice480P: cloneGroupValuePointer(source.VideoPrice480P),
VideoPrice720P: cloneGroupValuePointer(source.VideoPrice720P),
VideoPrice1080P: cloneGroupValuePointer(source.VideoPrice1080P),
VideoModelPrices: cloneGroupVideoModelPrices(source.VideoModelPrices),
WebSearchPricePerCall: cloneGroupValuePointer(source.WebSearchPricePerCall),
SearchPricePer1k: cloneGroupValuePointer(source.SearchPricePer1k),
AudioRealtimePricePerMin: cloneGroupValuePointer(source.AudioRealtimePricePerMin),
AudioTTSPricePerMillionChars: cloneGroupValuePointer(source.AudioTTSPricePerMillionChars),
AudioSTTPricePerHour: cloneGroupValuePointer(source.AudioSTTPricePerHour),
ClaudeCodeOnly: source.ClaudeCodeOnly,
FallbackGroupID: cloneGroupValuePointer(source.FallbackGroupID),
FallbackGroupIDOnInvalidRequest: cloneGroupValuePointer(source.FallbackGroupIDOnInvalidRequest),
ModelRouting: cloneGroupModelRouting(source.ModelRouting),
ModelRoutingEnabled: source.ModelRoutingEnabled,
MCPXMLInject: source.MCPXMLInject,
SupportedModelScopes: append([]string(nil), source.SupportedModelScopes...),
SortOrder: source.SortOrder,
AllowMessagesDispatch: source.AllowMessagesDispatch,
AllowLive: source.AllowLive,
RequireOAuthOnly: source.RequireOAuthOnly,
RequirePrivacySet: source.RequirePrivacySet,
DefaultMappedModel: source.DefaultMappedModel,
MessagesDispatchModelConfig: cloneGroupMessagesDispatchModelConfig(source.MessagesDispatchModelConfig),
ModelsListConfig: GroupModelsListConfig{
Enabled: source.ModelsListConfig.Enabled,
Models: append([]string(nil), source.ModelsListConfig.Models...),
},
RPMLimit: source.RPMLimit,
MaxReasoningEffort: source.MaxReasoningEffort,
ReasoningEffortMappings: append([]ReasoningEffortMapping(nil), source.ReasoningEffortMappings...),
}
}
// RecoverDuplicateGroup performs a read-only lookup for a copy that was already
// committed for the same actor, source group, and idempotency key.
func (s *adminServiceImpl) RecoverDuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error) {
operationID := duplicateGroupOperationID(id, actorScope, operationKey)
if operationID == "" {
return nil, nil
}
if s.groupDuplicateRepo == nil {
return nil, errors.New("group duplicate repository is not configured")
}
group, err := s.groupDuplicateRepo.FindByDuplicateOperationID(ctx, operationID)
if err != nil {
return nil, fmt.Errorf("find duplicate group operation: %w", err)
}
if group == nil {
return nil, nil
}
hydrated, err := s.groupRepo.GetByID(ctx, group.ID)
if err != nil {
return nil, fmt.Errorf("load recovered duplicate group: %w", err)
}
return hydrated, nil
}
// DuplicateGroup creates an inactive copy of a group's configuration and exact
// account priorities. The repository commits the group, bindings, and outbox
// event atomically so a failed binding never leaves an orphan group.
func (s *adminServiceImpl) DuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error) {
existing, err := s.RecoverDuplicateGroup(ctx, id, actorScope, operationKey)
if err != nil {
return nil, err
}
if existing != nil {
return existing, nil
}
source, err := s.groupRepo.GetByID(ctx, id)
if err != nil {
return nil, err
}
if s.groupDuplicateRepo == nil {
return nil, errors.New("group duplicate repository is not configured")
}
duplicate := cloneGroupForDuplicate(source, duplicateGroupOperationID(id, actorScope, operationKey))
sanitizeGroupReasoningEffortPolicy(duplicate)
for copyNumber := 1; ; copyNumber++ {
duplicate.Name = duplicateGroupName(source.Name, copyNumber)
duplicate.ID = 0
duplicate.CreatedAt = time.Time{}
duplicate.UpdatedAt = time.Time{}
if err := s.groupDuplicateRepo.CreateFromSource(ctx, duplicate, source.ID); err == nil {
hydrated, loadErr := s.groupRepo.GetByID(ctx, duplicate.ID)
if loadErr != nil {
return nil, fmt.Errorf("load duplicate group: %w", loadErr)
}
return hydrated, nil
} else if !errors.Is(err, ErrGroupExists) {
return nil, fmt.Errorf("create duplicate group: %w", err)
}
// A unique conflict can be either the generated name or the operation ID.
// Recover first; if no operation row exists, advance to the next name.
recovered, recoverErr := s.RecoverDuplicateGroup(ctx, id, actorScope, operationKey)
if recoverErr != nil {
return nil, recoverErr
}
if recovered != nil {
return recovered, nil
}
}
}
@@ -0,0 +1,292 @@
//go:build unit
package service
import (
"context"
"errors"
"strings"
"testing"
"time"
"unicode/utf8"
"github.com/stretchr/testify/require"
)
type duplicateGroupRepoStub struct {
GroupRepository
nextID int64
groups map[int64]*Group
names map[string]struct{}
byOperation map[string]int64
sourceBindings map[int64][]AccountGroup
createdBindings map[int64][]AccountGroup
createdFromSources []int64
atomicCreateErr error
}
func newDuplicateGroupRepoStub(source *Group) *duplicateGroupRepoStub {
repo := &duplicateGroupRepoStub{
nextID: 100,
groups: make(map[int64]*Group),
names: make(map[string]struct{}),
byOperation: make(map[string]int64),
sourceBindings: make(map[int64][]AccountGroup),
createdBindings: make(map[int64][]AccountGroup),
}
if source != nil {
repo.groups[source.ID] = source
repo.names[source.Name] = struct{}{}
}
return repo
}
func cloneGroupForDuplicateTest(group *Group) *Group {
if group == nil {
return nil
}
cloned := *group
cloned.DailyLimitUSD = cloneGroupValuePointer(group.DailyLimitUSD)
cloned.WeeklyLimitUSD = cloneGroupValuePointer(group.WeeklyLimitUSD)
cloned.MonthlyLimitUSD = cloneGroupValuePointer(group.MonthlyLimitUSD)
cloned.ImagePrice1K = cloneGroupValuePointer(group.ImagePrice1K)
cloned.ImagePrice2K = cloneGroupValuePointer(group.ImagePrice2K)
cloned.ImagePrice4K = cloneGroupValuePointer(group.ImagePrice4K)
cloned.VideoPrice480P = cloneGroupValuePointer(group.VideoPrice480P)
cloned.VideoPrice720P = cloneGroupValuePointer(group.VideoPrice720P)
cloned.VideoPrice1080P = cloneGroupValuePointer(group.VideoPrice1080P)
cloned.WebSearchPricePerCall = cloneGroupValuePointer(group.WebSearchPricePerCall)
cloned.FallbackGroupID = cloneGroupValuePointer(group.FallbackGroupID)
cloned.FallbackGroupIDOnInvalidRequest = cloneGroupValuePointer(group.FallbackGroupIDOnInvalidRequest)
cloned.ModelRouting = cloneGroupModelRouting(group.ModelRouting)
cloned.SupportedModelScopes = append([]string(nil), group.SupportedModelScopes...)
cloned.MessagesDispatchModelConfig = cloneGroupMessagesDispatchModelConfig(group.MessagesDispatchModelConfig)
cloned.ModelsListConfig.Models = append([]string(nil), group.ModelsListConfig.Models...)
cloned.AccountGroups = append([]AccountGroup(nil), group.AccountGroups...)
return &cloned
}
func (r *duplicateGroupRepoStub) GetByID(_ context.Context, id int64) (*Group, error) {
group := r.groups[id]
if group == nil {
return nil, ErrGroupNotFound
}
cloned := cloneGroupForDuplicateTest(group)
cloned.Hydrated = true
return cloned, nil
}
func (r *duplicateGroupRepoStub) FindByDuplicateOperationID(_ context.Context, operationID string) (*Group, error) {
id := r.byOperation[operationID]
if id == 0 {
return nil, nil
}
return cloneGroupForDuplicateTest(r.groups[id]), nil
}
func (r *duplicateGroupRepoStub) CreateFromSource(_ context.Context, group *Group, sourceGroupID int64) error {
if r.atomicCreateErr != nil {
return r.atomicCreateErr
}
if group.DuplicateOperationID != "" {
if _, exists := r.byOperation[group.DuplicateOperationID]; exists {
return ErrGroupExists
}
}
if _, exists := r.names[group.Name]; exists {
return ErrGroupExists
}
r.nextID++
group.ID = r.nextID
group.CreatedAt = time.Now().UTC()
group.UpdatedAt = group.CreatedAt
bindings := append([]AccountGroup(nil), r.sourceBindings[sourceGroupID]...)
for i := range bindings {
bindings[i].GroupID = group.ID
}
group.AccountCount = int64(len(bindings))
group.ActiveAccountCount = int64(len(bindings))
r.createdBindings[group.ID] = bindings
r.createdFromSources = append(r.createdFromSources, sourceGroupID)
r.names[group.Name] = struct{}{}
r.groups[group.ID] = cloneGroupForDuplicateTest(group)
if group.DuplicateOperationID != "" {
r.byOperation[group.DuplicateOperationID] = group.ID
}
return nil
}
func groupDuplicateTestPointer[T any](value T) *T { return &value }
func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing.T) {
createdAt := time.Date(2026, time.July, 1, 2, 3, 4, 0, time.UTC)
source := &Group{
ID: 41,
Name: "高级订阅",
Description: "configuration",
Platform: PlatformOpenAI,
RateMultiplier: 1.75,
PeakRateEnabled: true,
PeakStart: "09:00",
PeakEnd: "18:00",
PeakRateMultiplier: 1.2,
IsExclusive: true,
Status: StatusActive,
Hydrated: true,
SubscriptionType: SubscriptionTypeSubscription,
DailyLimitUSD: groupDuplicateTestPointer(11.0),
WeeklyLimitUSD: groupDuplicateTestPointer(22.0),
MonthlyLimitUSD: groupDuplicateTestPointer(33.0),
DefaultValidityDays: 91,
AllowImageGeneration: true,
AllowBatchImageGeneration: true,
ImageRateIndependent: true,
ImageRateMultiplier: 1.4,
ImagePrice1K: groupDuplicateTestPointer(0.01),
ImagePrice2K: groupDuplicateTestPointer(0.02),
ImagePrice4K: groupDuplicateTestPointer(0.04),
BatchImageDiscountMultiplier: 0.4,
BatchImageHoldMultiplier: 0.7,
VideoRateIndependent: true,
VideoRateMultiplier: 2.1,
VideoPrice480P: groupDuplicateTestPointer(0.1),
VideoPrice720P: groupDuplicateTestPointer(0.2),
VideoPrice1080P: groupDuplicateTestPointer(0.3),
VideoModelPrices: map[string]map[string]float64{
VideoPriceFamilyGrokImagineVideo15: {VideoBillingResolution720P: 0.14},
},
WebSearchPricePerCall: groupDuplicateTestPointer(0.005),
ClaudeCodeOnly: true,
FallbackGroupID: groupDuplicateTestPointer(int64(7)),
FallbackGroupIDOnInvalidRequest: groupDuplicateTestPointer(int64(8)),
ModelRouting: map[string][]int64{"gpt-*": {13, 17}},
ModelRoutingEnabled: true,
MCPXMLInject: true,
SupportedModelScopes: []string{"claude", "gemini_text"},
SortOrder: 9,
AllowMessagesDispatch: true,
AllowLive: true,
RequireOAuthOnly: true,
RequirePrivacySet: true,
DefaultMappedModel: "gpt-5.4",
MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4",
SonnetMappedModel: "gpt-5.3",
HaikuMappedModel: "gpt-5-mini",
ExactModelMappings: map[string]string{"claude-special": "gpt-special"},
},
ModelsListConfig: GroupModelsListConfig{Enabled: true, Models: []string{"gpt-5.4", "gpt-5-mini"}},
RPMLimit: 99,
MaxReasoningEffort: "medium",
ReasoningEffortMappings: []ReasoningEffortMapping{{From: "max", To: "xhigh"}},
CreatedAt: createdAt,
UpdatedAt: createdAt,
AccountCount: 12,
ActiveAccountCount: 8,
RateLimitedAccountCount: 2,
DuplicateOperationID: "old-operation-must-not-copy",
AccountGroups: []AccountGroup{{AccountID: 13, GroupID: 41, Priority: 37}},
}
repo := newDuplicateGroupRepoStub(source)
repo.sourceBindings[source.ID] = []AccountGroup{
{AccountID: 13, GroupID: source.ID, Priority: 37},
{AccountID: 17, GroupID: source.ID, Priority: 8},
}
svc := &adminServiceImpl{groupRepo: repo, groupDuplicateRepo: repo}
duplicate, err := svc.DuplicateGroup(context.Background(), source.ID, "admin:7", "stable-key")
require.NoError(t, err)
require.NotEqual(t, source.ID, duplicate.ID)
require.Equal(t, "高级订阅 (Copy)", duplicate.Name)
require.Equal(t, duplicateGroupInactiveStatus, duplicate.Status)
require.True(t, duplicate.Hydrated, "the duplicate response is reloaded with derived counts")
require.Equal(t, source.Description, duplicate.Description)
require.Equal(t, source.Platform, duplicate.Platform)
require.Equal(t, source.RateMultiplier, duplicate.RateMultiplier)
require.Equal(t, source.PeakRateMultiplier, duplicate.PeakRateMultiplier)
require.Equal(t, source.DefaultValidityDays, duplicate.DefaultValidityDays)
require.Equal(t, source.ImagePrice4K, duplicate.ImagePrice4K)
require.Equal(t, source.VideoModelPrices, duplicate.VideoModelPrices)
require.Equal(t, source.WebSearchPricePerCall, duplicate.WebSearchPricePerCall)
require.Equal(t, source.FallbackGroupID, duplicate.FallbackGroupID)
require.Equal(t, source.ModelRouting, duplicate.ModelRouting)
require.Equal(t, source.MessagesDispatchModelConfig, duplicate.MessagesDispatchModelConfig)
require.Equal(t, source.ModelsListConfig, duplicate.ModelsListConfig)
require.Equal(t, source.RPMLimit, duplicate.RPMLimit)
require.Equal(t, source.MaxReasoningEffort, duplicate.MaxReasoningEffort)
require.Equal(t, source.ReasoningEffortMappings, duplicate.ReasoningEffortMappings)
require.EqualValues(t, 2, duplicate.AccountCount)
require.EqualValues(t, 2, duplicate.ActiveAccountCount)
require.NotEmpty(t, duplicate.DuplicateOperationID)
require.Equal(t, []int64{source.ID}, repo.createdFromSources)
require.Equal(t, []AccountGroup{
{AccountID: 13, GroupID: duplicate.ID, Priority: 37},
{AccountID: 17, GroupID: duplicate.ID, Priority: 8},
}, repo.createdBindings[duplicate.ID])
duplicate.ModelRouting["gpt-*"][0] = 999
duplicate.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P] = 999
duplicate.SupportedModelScopes[0] = "changed"
duplicate.MessagesDispatchModelConfig.ExactModelMappings["claude-special"] = "changed"
duplicate.ModelsListConfig.Models[0] = "changed"
duplicate.ReasoningEffortMappings[0].To = "changed"
*duplicate.DailyLimitUSD = 999
require.Equal(t, int64(13), source.ModelRouting["gpt-*"][0])
require.Equal(t, 0.14, source.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P])
require.Equal(t, "claude", source.SupportedModelScopes[0])
require.Equal(t, "gpt-special", source.MessagesDispatchModelConfig.ExactModelMappings["claude-special"])
require.Equal(t, "gpt-5.4", source.ModelsListConfig.Models[0])
require.Equal(t, "xhigh", source.ReasoningEffortMappings[0].To)
require.Equal(t, 11.0, *source.DailyLimitUSD)
}
func TestDuplicateGroupRecoversSameOperationAndScopesByAdmin(t *testing.T) {
source := &Group{ID: 9, Name: "team", Platform: PlatformAnthropic, Status: StatusActive}
repo := newDuplicateGroupRepoStub(source)
svc := &adminServiceImpl{groupRepo: repo, groupDuplicateRepo: repo}
ctx := context.Background()
first, err := svc.DuplicateGroup(ctx, source.ID, "admin:7", "same-key")
require.NoError(t, err)
retry, err := svc.DuplicateGroup(ctx, source.ID, "admin:7", "same-key")
require.NoError(t, err)
recovered, err := svc.RecoverDuplicateGroup(ctx, source.ID, "admin:7", "same-key")
require.NoError(t, err)
otherAdmin, err := svc.DuplicateGroup(ctx, source.ID, "admin:8", "same-key")
require.NoError(t, err)
require.Equal(t, first.ID, retry.ID)
require.Equal(t, first.ID, recovered.ID)
require.NotEqual(t, first.ID, otherAdmin.ID)
require.Equal(t, "team (Copy 2)", otherAdmin.Name)
}
func TestDuplicateGroupAdvancesNameAndTruncatesUnicodeByRunes(t *testing.T) {
source := &Group{ID: 12, Name: "team", Platform: PlatformAnthropic, Status: StatusActive}
repo := newDuplicateGroupRepoStub(source)
repo.names["team (Copy)"] = struct{}{}
svc := &adminServiceImpl{groupRepo: repo, groupDuplicateRepo: repo}
duplicate, err := svc.DuplicateGroup(context.Background(), source.ID, "admin:1", "")
require.NoError(t, err)
require.Equal(t, "team (Copy 2)", duplicate.Name)
unicodeName := duplicateGroupName(strings.Repeat("组", 100), 23)
require.Equal(t, maxGroupNameRunes, utf8.RuneCountInString(unicodeName))
require.True(t, strings.HasSuffix(unicodeName, " (Copy 23)"))
}
func TestDuplicateGroupAtomicCreateFailureReturnsNoCopy(t *testing.T) {
source := &Group{ID: 15, Name: "team", Platform: PlatformAnthropic, Status: StatusActive}
repo := newDuplicateGroupRepoStub(source)
repo.atomicCreateErr = errors.New("binding insert failed")
svc := &adminServiceImpl{groupRepo: repo, groupDuplicateRepo: repo}
duplicate, err := svc.DuplicateGroup(context.Background(), source.ID, "admin:1", "key")
require.ErrorContains(t, err, "binding insert failed")
require.Nil(t, duplicate)
require.Len(t, repo.groups, 1)
require.Empty(t, repo.byOperation)
}
@@ -0,0 +1,87 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
// groupPlatformRepoStub 只实现 UpdateGroup 走到的两个方法,其余靠内嵌接口占位。
type groupPlatformRepoStub struct {
GroupRepository
group *Group
updated *Group
}
func (r *groupPlatformRepoStub) GetByID(_ context.Context, _ int64) (*Group, error) {
cloned := *r.group
return &cloned, nil
}
func (r *groupPlatformRepoStub) Update(_ context.Context, group *Group) error {
r.updated = group
return nil
}
type channelCacheInvalidatorSpy struct {
calls int
}
func (s *channelCacheInvalidatorSpy) InvalidateCache() { s.calls++ }
// 渠道缓存持有 groupID → platform,而渠道定价/模型映射/模型白名单都按平台严格隔离。
// 改了分组平台却不失效缓存,最长 10 分钟内这些查找仍按旧平台匹配(静默走错价)。
func TestUpdateGroupInvalidatesChannelCacheOnPlatformChange(t *testing.T) {
tests := []struct {
name string
fromPlatform string
inputPlatform string
wantCalls int
}{
{
name: "platform changed invalidates",
fromPlatform: PlatformAnthropic,
inputPlatform: PlatformOpenAI,
wantCalls: 1,
},
{
name: "same platform does not invalidate",
fromPlatform: PlatformAnthropic,
inputPlatform: PlatformAnthropic,
wantCalls: 0,
},
{
// 请求里不带 platform 字段时不应该动缓存
name: "platform omitted does not invalidate",
fromPlatform: PlatformAnthropic,
inputPlatform: "",
wantCalls: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: tt.fromPlatform}}
spy := &channelCacheInvalidatorSpy{}
svc := &adminServiceImpl{groupRepo: repo, channelCacheInvalidator: spy}
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: tt.inputPlatform})
require.NoError(t, err)
require.NotNil(t, got)
require.Equal(t, tt.wantCalls, spy.calls)
})
}
}
// 依赖可以不注入(例如测试或裁剪构建),此时不应 panic——缓存靠 TTL 自然重建。
func TestUpdateGroupWithoutChannelCacheInvalidator(t *testing.T) {
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: PlatformAnthropic}}
svc := &adminServiceImpl{groupRepo: repo}
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: PlatformOpenAI})
require.NoError(t, err)
require.Equal(t, PlatformOpenAI, got.Platform)
}
+608
View File
@@ -0,0 +1,608 @@
package service
import (
"context"
"fmt"
"io"
"net/http"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/util/httputil"
)
// Proxy management implementations
func (s *adminServiceImpl) ListProxies(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]Proxy, int64, error) {
params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder}
proxies, result, err := s.proxyRepo.ListWithFilters(ctx, params, protocol, status, search)
if err != nil {
return nil, 0, err
}
return proxies, result.Total, nil
}
func (s *adminServiceImpl) ListProxiesWithAccountCount(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]ProxyWithAccountCount, int64, error) {
params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder}
proxies, result, err := s.proxyRepo.ListWithFiltersAndAccountCount(ctx, params, protocol, status, search)
if err != nil {
return nil, 0, err
}
s.attachProxyLatency(ctx, proxies)
return proxies, result.Total, nil
}
func (s *adminServiceImpl) GetAllProxies(ctx context.Context) ([]Proxy, error) {
return s.proxyRepo.ListActive(ctx)
}
func (s *adminServiceImpl) GetAllProxiesWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error) {
proxies, err := s.proxyRepo.ListActiveWithAccountCount(ctx)
if err != nil {
return nil, err
}
s.attachProxyLatency(ctx, proxies)
return proxies, nil
}
func (s *adminServiceImpl) GetProxy(ctx context.Context, id int64) (*Proxy, error) {
return s.proxyRepo.GetByID(ctx, id)
}
func (s *adminServiceImpl) GetProxiesByIDs(ctx context.Context, ids []int64) ([]Proxy, error) {
return s.proxyRepo.ListByIDs(ctx, ids)
}
func (s *adminServiceImpl) CreateProxy(ctx context.Context, input *CreateProxyInput) (*Proxy, error) {
// 规范化 fallback_mode
mode := input.FallbackMode
if mode == "" {
mode = FallbackModeNone
}
// 校验:mode=proxy 必须有 backup
if mode == FallbackModeProxy && input.BackupProxyID == nil {
return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy")
}
if input.ExpiryWarnDays < 0 {
return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0")
}
proxy := &Proxy{
Name: input.Name,
Protocol: input.Protocol,
Host: input.Host,
Port: input.Port,
Username: input.Username,
Password: input.Password,
Status: StatusActive,
ExpiresAt: input.ExpiresAt,
FallbackMode: mode,
BackupProxyID: input.BackupProxyID,
ExpiryWarnDays: input.ExpiryWarnDays,
}
if err := s.proxyRepo.Create(ctx, proxy); err != nil {
return nil, err
}
// Probe latency asynchronously so creation isn't blocked by network timeout.
go s.probeProxyLatency(context.Background(), proxy)
return proxy, nil
}
func (s *adminServiceImpl) UpdateProxy(ctx context.Context, id int64, input *UpdateProxyInput) (*Proxy, error) {
// 校验:backup_proxy_id 不能是自身
if input.BackupProxyID != nil && *input.BackupProxyID == id {
return nil, infraerrors.BadRequest("PROXY_BACKUP_SELF", "backup proxy cannot be itself")
}
// 规范化 fallback_mode
mode := input.FallbackMode
if mode == "" {
mode = FallbackModeNone
}
// 校验:mode=proxy 必须有 backup
if mode == FallbackModeProxy && input.BackupProxyID == nil {
return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy")
}
if input.ExpiryWarnDays < 0 {
return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0")
}
proxy, err := s.proxyRepo.GetByID(ctx, id)
if err != nil {
return nil, err
}
if input.Name != "" {
proxy.Name = input.Name
}
if input.Protocol != "" {
proxy.Protocol = input.Protocol
}
if input.Host != "" {
proxy.Host = input.Host
}
if input.Port != 0 {
proxy.Port = input.Port
}
if input.Username != "" {
proxy.Username = input.Username
}
if input.Password != "" {
proxy.Password = input.Password
}
if input.Status != "" {
proxy.Status = input.Status
}
// 透传有效期与回退字段
proxy.ExpiresAt = input.ExpiresAt
proxy.FallbackMode = mode
proxy.BackupProxyID = input.BackupProxyID
proxy.ExpiryWarnDays = input.ExpiryWarnDays
if err := s.proxyRepo.Update(ctx, proxy); err != nil {
return nil, err
}
return proxy, nil
}
func (s *adminServiceImpl) DeleteProxy(ctx context.Context, id int64) error {
count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id)
if err != nil {
return err
}
if count > 0 {
return ErrProxyInUse
}
return s.proxyRepo.Delete(ctx, id)
}
func (s *adminServiceImpl) BatchDeleteProxies(ctx context.Context, ids []int64) (*ProxyBatchDeleteResult, error) {
result := &ProxyBatchDeleteResult{}
if len(ids) == 0 {
return result, nil
}
for _, id := range ids {
count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id)
if err != nil {
result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{
ID: id,
Reason: err.Error(),
})
continue
}
if count > 0 {
result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{
ID: id,
Reason: ErrProxyInUse.Error(),
})
continue
}
if err := s.proxyRepo.Delete(ctx, id); err != nil {
result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{
ID: id,
Reason: err.Error(),
})
continue
}
result.DeletedIDs = append(result.DeletedIDs, id)
}
return result, nil
}
func (s *adminServiceImpl) GetProxyAccounts(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error) {
return s.proxyRepo.ListAccountSummariesByProxyID(ctx, proxyID)
}
func (s *adminServiceImpl) CheckProxyExists(ctx context.Context, host string, port int, username, password string) (bool, error) {
return s.proxyRepo.ExistsByHostPortAuth(ctx, host, port, username, password)
}
func (s *adminServiceImpl) TestProxy(ctx context.Context, id int64) (*ProxyTestResult, error) {
proxy, err := s.proxyRepo.GetByID(ctx, id)
if err != nil {
return nil, err
}
proxyURL := proxy.URL()
exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL)
if err != nil {
s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{
Success: false,
Message: err.Error(),
UpdatedAt: time.Now(),
})
return &ProxyTestResult{
Success: false,
Message: err.Error(),
}, nil
}
latency := latencyMs
s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{
Success: true,
LatencyMs: &latency,
Message: "Proxy is accessible",
IPAddress: exitInfo.IP,
Country: exitInfo.Country,
CountryCode: exitInfo.CountryCode,
Region: exitInfo.Region,
City: exitInfo.City,
UpdatedAt: time.Now(),
})
return &ProxyTestResult{
Success: true,
Message: "Proxy is accessible",
LatencyMs: latencyMs,
IPAddress: exitInfo.IP,
City: exitInfo.City,
Region: exitInfo.Region,
Country: exitInfo.Country,
CountryCode: exitInfo.CountryCode,
}, nil
}
func (s *adminServiceImpl) CheckProxyQuality(ctx context.Context, id int64) (*ProxyQualityCheckResult, error) {
proxy, err := s.proxyRepo.GetByID(ctx, id)
if err != nil {
return nil, err
}
result := &ProxyQualityCheckResult{
ProxyID: id,
Score: 100,
Grade: "A",
CheckedAt: time.Now().Unix(),
Items: make([]ProxyQualityCheckItem, 0, len(proxyQualityTargets)+1),
}
proxyURL := proxy.URL()
if s.proxyProber == nil {
result.Items = append(result.Items, ProxyQualityCheckItem{
Target: "base_connectivity",
Status: "fail",
Message: "代理探测服务未配置",
})
result.FailedCount++
finalizeProxyQualityResult(result)
s.saveProxyQualitySnapshot(ctx, id, result, nil)
return result, nil
}
exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL)
if err != nil {
result.Items = append(result.Items, ProxyQualityCheckItem{
Target: "base_connectivity",
Status: "fail",
LatencyMs: latencyMs,
Message: err.Error(),
})
result.FailedCount++
finalizeProxyQualityResult(result)
s.saveProxyQualitySnapshot(ctx, id, result, nil)
return result, nil
}
result.ExitIP = exitInfo.IP
result.Country = exitInfo.Country
result.CountryCode = exitInfo.CountryCode
result.BaseLatencyMs = latencyMs
result.Items = append(result.Items, ProxyQualityCheckItem{
Target: "base_connectivity",
Status: "pass",
LatencyMs: latencyMs,
Message: "代理出口连通正常",
})
result.PassedCount++
client, err := httpclient.GetClient(httpclient.Options{
ProxyURL: proxyURL,
Timeout: proxyQualityRequestTimeout,
ResponseHeaderTimeout: proxyQualityResponseHeaderTimeout,
})
if err != nil {
result.Items = append(result.Items, ProxyQualityCheckItem{
Target: "http_client",
Status: "fail",
Message: fmt.Sprintf("创建检测客户端失败: %v", err),
})
result.FailedCount++
finalizeProxyQualityResult(result)
s.saveProxyQualitySnapshot(ctx, id, result, exitInfo)
return result, nil
}
for _, target := range proxyQualityTargets {
item := runProxyQualityTarget(ctx, client, target)
result.Items = append(result.Items, item)
switch item.Status {
case "pass":
result.PassedCount++
case "warn":
result.WarnCount++
case "challenge":
result.ChallengeCount++
default:
result.FailedCount++
}
}
finalizeProxyQualityResult(result)
s.saveProxyQualitySnapshot(ctx, id, result, exitInfo)
return result, nil
}
func runProxyQualityTarget(ctx context.Context, client *http.Client, target proxyQualityTarget) ProxyQualityCheckItem {
item := ProxyQualityCheckItem{
Target: target.Target,
}
req, err := http.NewRequestWithContext(ctx, target.Method, target.URL, nil)
if err != nil {
item.Status = "fail"
item.Message = fmt.Sprintf("构建请求失败: %v", err)
return item
}
req.Header.Set("Accept", "application/json,text/html,*/*")
req.Header.Set("User-Agent", proxyQualityClientUserAgent)
start := time.Now()
resp, err := client.Do(req)
if err != nil {
item.Status = "fail"
item.LatencyMs = time.Since(start).Milliseconds()
item.Message = fmt.Sprintf("请求失败: %v", err)
return item
}
defer func() { _ = resp.Body.Close() }()
item.LatencyMs = time.Since(start).Milliseconds()
item.HTTPStatus = resp.StatusCode
body, readErr := io.ReadAll(io.LimitReader(resp.Body, proxyQualityMaxBodyBytes+1))
if readErr != nil {
item.Status = "fail"
item.Message = fmt.Sprintf("读取响应失败: %v", readErr)
return item
}
if int64(len(body)) > proxyQualityMaxBodyBytes {
body = body[:proxyQualityMaxBodyBytes]
}
// Cloudflare challenge 检测
if httputil.IsCloudflareChallengeResponse(resp.StatusCode, resp.Header, body) {
item.Status = "challenge"
item.CFRay = httputil.ExtractCloudflareRayID(resp.Header, body)
item.Message = "命中 Cloudflare challenge"
return item
}
if _, ok := target.AllowedStatuses[resp.StatusCode]; ok {
// 白名单内的状态码均代表目标可达:2xx 表示接口直接可用,
// 401/405 等是无鉴权探测的预期结果,同样视为连通正常,不再扣分。
item.Status = "pass"
if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices {
item.Message = fmt.Sprintf("HTTP %d", resp.StatusCode)
} else {
item.Message = fmt.Sprintf("HTTP %d(目标可达)", resp.StatusCode)
}
return item
}
if resp.StatusCode == http.StatusTooManyRequests {
item.Status = "warn"
item.Message = "目标返回 429,可能存在频控"
return item
}
item.Status = "fail"
item.Message = fmt.Sprintf("非预期状态码: %d", resp.StatusCode)
return item
}
func finalizeProxyQualityResult(result *ProxyQualityCheckResult) {
if result == nil {
return
}
score := 100 - result.WarnCount*10 - result.FailedCount*22 - result.ChallengeCount*30
if score < 0 {
score = 0
}
result.Score = score
result.Grade = proxyQualityGrade(score)
result.Summary = fmt.Sprintf(
"通过 %d 项,告警 %d 项,失败 %d 项,挑战 %d 项",
result.PassedCount,
result.WarnCount,
result.FailedCount,
result.ChallengeCount,
)
}
func proxyQualityGrade(score int) string {
switch {
case score >= 90:
return "A"
case score >= 75:
return "B"
case score >= 60:
return "C"
case score >= 40:
return "D"
default:
return "F"
}
}
func proxyQualityOverallStatus(result *ProxyQualityCheckResult) string {
if result == nil {
return ""
}
if result.ChallengeCount > 0 {
return "challenge"
}
if result.FailedCount > 0 {
return "failed"
}
if result.WarnCount > 0 {
return "warn"
}
if result.PassedCount > 0 {
return "healthy"
}
return "failed"
}
func proxyQualityFirstCFRay(result *ProxyQualityCheckResult) string {
if result == nil {
return ""
}
for _, item := range result.Items {
if item.CFRay != "" {
return item.CFRay
}
}
return ""
}
func proxyQualityBaseConnectivityPass(result *ProxyQualityCheckResult) bool {
if result == nil {
return false
}
for _, item := range result.Items {
if item.Target == "base_connectivity" {
return item.Status == "pass"
}
}
return false
}
func (s *adminServiceImpl) saveProxyQualitySnapshot(ctx context.Context, proxyID int64, result *ProxyQualityCheckResult, exitInfo *ProxyExitInfo) {
if result == nil {
return
}
score := result.Score
checkedAt := result.CheckedAt
info := &ProxyLatencyInfo{
Success: proxyQualityBaseConnectivityPass(result),
Message: result.Summary,
QualityStatus: proxyQualityOverallStatus(result),
QualityScore: &score,
QualityGrade: result.Grade,
QualitySummary: result.Summary,
QualityCheckedAt: &checkedAt,
QualityCFRay: proxyQualityFirstCFRay(result),
UpdatedAt: time.Now(),
}
if result.BaseLatencyMs > 0 {
latency := result.BaseLatencyMs
info.LatencyMs = &latency
}
if exitInfo != nil {
info.IPAddress = exitInfo.IP
info.Country = exitInfo.Country
info.CountryCode = exitInfo.CountryCode
info.Region = exitInfo.Region
info.City = exitInfo.City
}
s.saveProxyLatency(ctx, proxyID, info)
}
func (s *adminServiceImpl) probeProxyLatency(ctx context.Context, proxy *Proxy) {
if s.proxyProber == nil || proxy == nil {
return
}
exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxy.URL())
if err != nil {
s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{
Success: false,
Message: err.Error(),
UpdatedAt: time.Now(),
})
return
}
latency := latencyMs
s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{
Success: true,
LatencyMs: &latency,
Message: "Proxy is accessible",
IPAddress: exitInfo.IP,
Country: exitInfo.Country,
CountryCode: exitInfo.CountryCode,
Region: exitInfo.Region,
City: exitInfo.City,
UpdatedAt: time.Now(),
})
}
func (s *adminServiceImpl) attachProxyLatency(ctx context.Context, proxies []ProxyWithAccountCount) {
if s.proxyLatencyCache == nil || len(proxies) == 0 {
return
}
ids := make([]int64, 0, len(proxies))
for i := range proxies {
ids = append(ids, proxies[i].ID)
}
latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, ids)
if err != nil {
logger.LegacyPrintf("service.admin", "Warning: load proxy latency cache failed: %v", err)
return
}
for i := range proxies {
info := latencies[proxies[i].ID]
if info == nil {
continue
}
if info.Success {
proxies[i].LatencyStatus = "success"
proxies[i].LatencyMs = info.LatencyMs
} else {
proxies[i].LatencyStatus = "failed"
}
proxies[i].LatencyMessage = info.Message
proxies[i].IPAddress = info.IPAddress
proxies[i].Country = info.Country
proxies[i].CountryCode = info.CountryCode
proxies[i].Region = info.Region
proxies[i].City = info.City
proxies[i].QualityStatus = info.QualityStatus
proxies[i].QualityScore = info.QualityScore
proxies[i].QualityGrade = info.QualityGrade
proxies[i].QualitySummary = info.QualitySummary
proxies[i].QualityChecked = info.QualityCheckedAt
}
}
func (s *adminServiceImpl) saveProxyLatency(ctx context.Context, proxyID int64, info *ProxyLatencyInfo) {
if s.proxyLatencyCache == nil || info == nil {
return
}
merged := *info
if latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, []int64{proxyID}); err == nil {
if existing := latencies[proxyID]; existing != nil {
if merged.QualityCheckedAt == nil &&
merged.QualityScore == nil &&
merged.QualityGrade == "" &&
merged.QualityStatus == "" &&
merged.QualitySummary == "" &&
merged.QualityCFRay == "" {
merged.QualityStatus = existing.QualityStatus
merged.QualityScore = existing.QualityScore
merged.QualityGrade = existing.QualityGrade
merged.QualitySummary = existing.QualitySummary
merged.QualityCheckedAt = existing.QualityCheckedAt
merged.QualityCFRay = existing.QualityCFRay
}
}
}
if err := s.proxyLatencyCache.SetProxyLatency(ctx, proxyID, &merged); err != nil {
logger.LegacyPrintf("service.admin", "Warning: store proxy latency cache failed: %v", err)
}
}
+741
View File
@@ -0,0 +1,741 @@
package service
import (
"context"
"net/http"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
// AdminService interface defines admin management operations
type AdminService interface {
// User management
ListUsers(ctx context.Context, page, pageSize int, filters UserListFilters, sortBy, sortOrder string) ([]User, int64, error)
GetUser(ctx context.Context, id int64) (*User, error)
GetUserIncludeDeleted(ctx context.Context, id int64) (*User, error)
CreateUser(ctx context.Context, input *CreateUserInput) (*User, error)
UpdateUser(ctx context.Context, id int64, input *UpdateUserInput) (*User, error)
DeleteUser(ctx context.Context, id int64) error
UpdateUserBalance(ctx context.Context, userID int64, balance float64, operation string, notes string) (*User, error)
BatchUpdateConcurrency(ctx context.Context, userIDs []int64, value int, mode string) (int, error)
BatchUpdateLimits(ctx context.Context, userIDs []int64, concurrency, rpmLimit *int) (int, error)
GetUserAPIKeys(ctx context.Context, userID int64, page, pageSize int, sortBy, sortOrder string) ([]APIKey, int64, error)
GetUserUsageStats(ctx context.Context, userID int64, period string) (any, error)
GetUserRPMStatus(ctx context.Context, userID int64) (*UserRPMStatus, error)
// GetUserBalanceHistory returns paginated balance/concurrency change records for a user.
// codeType is optional - pass empty string to return all types.
// Also returns totalRecharged (sum of all positive balance top-ups).
GetUserBalanceHistory(ctx context.Context, userID int64, page, pageSize int, codeType string) ([]RedeemCode, int64, float64, error)
BindUserAuthIdentity(ctx context.Context, userID int64, input AdminBindAuthIdentityInput) (*AdminBoundAuthIdentity, error)
// Group management
ListGroups(ctx context.Context, page, pageSize int, platform, status, search string, isExclusive *bool, sortBy, sortOrder string) ([]Group, int64, error)
GetAllGroups(ctx context.Context) ([]Group, error)
GetAllGroupsByPlatform(ctx context.Context, platform string) ([]Group, error)
// GetAllGroupsIncludingInactive returns all groups regardless of status (active + disabled),
// ordered by sort_order then id. Used by the API Key group filter dropdown.
GetAllGroupsIncludingInactive(ctx context.Context) ([]Group, error)
GetGroup(ctx context.Context, id int64) (*Group, error)
GetGroupModelsListCandidates(ctx context.Context, id int64, platform string) ([]string, error)
CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error)
// DuplicateGroup creates an inactive independent copy of a group's configuration
// and account bindings while preserving each binding's priority.
DuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error)
// RecoverDuplicateGroup returns a previously committed copy for an ambiguous retry.
// It never creates a group.
RecoverDuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error)
UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error)
DeleteGroup(ctx context.Context, id int64) error
ListCompositeRoutes(ctx context.Context, groupID int64) ([]CompositeModelRoute, error)
CreateCompositeRoute(ctx context.Context, groupID int64, input CompositeRouteInput) (*CompositeModelRoute, error)
UpdateCompositeRoute(ctx context.Context, groupID, routeID int64, input CompositeRouteInput) (*CompositeModelRoute, error)
DeleteCompositeRoute(ctx context.Context, groupID, routeID int64) error
PreviewCompositeRoute(ctx context.Context, groupID int64, input CompositeRoutePreviewRequest) (*CompositeRouteDecision, error)
GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error)
GetGroupRateMultipliers(ctx context.Context, groupID int64) ([]UserGroupRateEntry, error)
ClearGroupRateMultipliers(ctx context.Context, groupID int64) error
BatchSetGroupRateMultipliers(ctx context.Context, groupID int64, entries []GroupRateMultiplierInput) error
ClearGroupRPMOverrides(ctx context.Context, groupID int64) error
BatchSetGroupRPMOverrides(ctx context.Context, groupID int64, entries []GroupRPMOverrideInput) error
UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error
// API Key management (admin)
AdminUpdateAPIKeyGroupID(ctx context.Context, keyID int64, groupID *int64) (*AdminUpdateAPIKeyGroupIDResult, error)
AdminResetAPIKeyRateLimitUsage(ctx context.Context, keyID int64) (*APIKey, error)
// ReplaceUserGroup 替换用户的专属分组:授予新分组权限、迁移 Key、移除旧分组权限
ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error)
// Account management
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]Account, int64, error)
// ListAccountsForSchedulerScoreFilter 返回符合过滤条件的全部账号(不分页),
// 作为账号列表页计算 OpenAI 调度分数的过滤范围池。
ListAccountsForSchedulerScoreFilter(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error)
// ListOpenAISchedulableAccountsForSchedulerScore 返回指定分组(nil 为未分组)内
// 可调度的 OpenAI 账号,用于按组计算调度分数。
ListOpenAISchedulableAccountsForSchedulerScore(ctx context.Context, groupID *int64) ([]Account, error)
GetAccount(ctx context.Context, id int64) (*Account, error)
GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error)
CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error)
// DuplicateAccount creates an independent account from an existing account's configuration.
// First-class runtime columns are intentionally reset by the normal account creation path.
DuplicateAccount(ctx context.Context, id int64, actorScope, operationKey string) (*Account, error)
// RecoverDuplicateAccount returns a previously committed duplicate for an ambiguous retry.
// It never creates an account.
RecoverDuplicateAccount(ctx context.Context, id int64, actorScope, operationKey string) (*Account, error)
UpdateAccount(ctx context.Context, id int64, input *UpdateAccountInput) (*Account, error)
// UpdateAccountExtra 仅对 Extra 做 JSONB 增量合并(key 级覆盖),不会影响其它字段或运行态键。
// 用于刷新流程持久化 account_uuid / org_uuid 等少量键,避免被全量快照覆盖。
UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error
DeleteAccount(ctx context.Context, id int64) error
RefreshAccountCredentials(ctx context.Context, id int64) (*Account, error)
ClearAccountError(ctx context.Context, id int64) (*Account, error)
SetAccountError(ctx context.Context, id int64, errorMsg string) error
// EnsureOpenAIPrivacy 检查 OpenAI OAuth 账号 privacy_mode,未设置则尝试关闭训练数据共享并持久化。
EnsureOpenAIPrivacy(ctx context.Context, account *Account) string
// EnsureAntigravityPrivacy 检查 Antigravity OAuth 账号 privacy_mode,未设置则调用 setUserSettings 并持久化。
EnsureAntigravityPrivacy(ctx context.Context, account *Account) string
// ForceOpenAIPrivacy 强制重新设置 OpenAI OAuth 账号隐私,无论当前状态。
ForceOpenAIPrivacy(ctx context.Context, account *Account) string
// ForceAntigravityPrivacy 强制重新设置 Antigravity OAuth 账号隐私,无论当前状态。
ForceAntigravityPrivacy(ctx context.Context, account *Account) string
SetAccountSchedulable(ctx context.Context, id int64, schedulable bool) (*Account, error)
BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error)
CheckMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error
// RevertAccountProxyFallback 将账号的 proxy_id 切回 proxy_fallback_origin_id,并清空 origin 字段。
// 若账号不存在返回 ErrAccountNotFound;若账号存在但不在 fallback 状态,返回 ErrAccountNotInFallback。
RevertAccountProxyFallback(ctx context.Context, id int64) error
// CreateShadow 为指定 OpenAI OAuth 母账号创建 spark 维度影子账号(一母一影)。
// 影子账号不持凭据(Credentials 恒为空),透传母账号凭据;继承母账号的 ProxyID。
CreateShadow(ctx context.Context, parentID int64, opts ShadowOptions) (*Account, error)
// Proxy management
ListProxies(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]Proxy, int64, error)
ListProxiesWithAccountCount(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]ProxyWithAccountCount, int64, error)
GetAllProxies(ctx context.Context) ([]Proxy, error)
GetAllProxiesWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error)
GetProxy(ctx context.Context, id int64) (*Proxy, error)
GetProxiesByIDs(ctx context.Context, ids []int64) ([]Proxy, error)
CreateProxy(ctx context.Context, input *CreateProxyInput) (*Proxy, error)
UpdateProxy(ctx context.Context, id int64, input *UpdateProxyInput) (*Proxy, error)
DeleteProxy(ctx context.Context, id int64) error
BatchDeleteProxies(ctx context.Context, ids []int64) (*ProxyBatchDeleteResult, error)
GetProxyAccounts(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error)
CheckProxyExists(ctx context.Context, host string, port int, username, password string) (bool, error)
TestProxy(ctx context.Context, id int64) (*ProxyTestResult, error)
CheckProxyQuality(ctx context.Context, id int64) (*ProxyQualityCheckResult, error)
// Redeem code management
ListRedeemCodes(ctx context.Context, page, pageSize int, codeType, status, search string, sortBy, sortOrder string) ([]RedeemCode, int64, error)
GetRedeemCode(ctx context.Context, id int64) (*RedeemCode, error)
GenerateRedeemCodes(ctx context.Context, input *GenerateRedeemCodesInput) ([]RedeemCode, error)
DeleteRedeemCode(ctx context.Context, id int64) error
BatchDeleteRedeemCodes(ctx context.Context, ids []int64) (int64, error)
ExpireRedeemCode(ctx context.Context, id int64) (*RedeemCode, error)
ResetAccountQuota(ctx context.Context, id int64) error
}
// CreateUserInput represents input for creating a new user via admin operations.
type CreateUserInput struct {
Email string
Password string
Username string
Notes string
Role string // 空字符串表示使用默认角色(user);合法值 admin/user
Balance *float64
Concurrency int
RPMLimit int
AllowedGroups []int64
// ActorAdminID 执行本次操作的管理员ID(来自JWT),仅用于权限敏感操作的审计日志。
ActorAdminID int64
}
type UpdateUserInput struct {
Email string
Password string
Username *string
Notes *string
Role string // 空字符串表示"未提供"(不修改);合法值 admin/user
Balance *float64 // 使用指针区分"未提供"和"设置为0"
Concurrency *int // 使用指针区分"未提供"和"设置为0"
RPMLimit *int // 使用指针区分"未提供"和"设置为0"
Status string
AllowedGroups *[]int64 // 使用指针区分"未提供"和"设置为空数组"
// GroupRates 用户专属分组倍率配置
// map[groupID]*ratenil 表示删除该分组的专属倍率
GroupRates map[int64]*float64
// ActorAdminID 执行本次操作的管理员ID(来自JWT),仅用于权限敏感操作的审计日志。
ActorAdminID int64
}
type AdminBindAuthIdentityInput struct {
ProviderType string
ProviderKey string
ProviderSubject string
Issuer *string
Metadata map[string]any
Channel *AdminBindAuthIdentityChannelInput
}
type AdminBindAuthIdentityChannelInput struct {
Channel string
ChannelAppID string
ChannelSubject string
Metadata map[string]any
}
type AdminBoundAuthIdentity struct {
UserID int64 `json:"user_id"`
ProviderType string `json:"provider_type"`
ProviderKey string `json:"provider_key"`
ProviderSubject string `json:"provider_subject"`
VerifiedAt *time.Time `json:"verified_at,omitempty"`
Issuer *string `json:"issuer,omitempty"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
Channel *AdminBoundAuthIdentityChannel `json:"channel,omitempty"`
}
type AdminBoundAuthIdentityChannel struct {
Channel string `json:"channel"`
ChannelAppID string `json:"channel_app_id"`
ChannelSubject string `json:"channel_subject"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type CreateGroupInput struct {
Name string
Description string
Platform string
RateMultiplier float64
IsExclusive bool
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
LongContextPricingEnabled bool
ModelPricing []ChannelModelPricing
// 图片生成计费配置(仅 antigravity 平台使用)
AllowImageGeneration bool
AllowBatchImageGeneration bool
ImageRateIndependent bool
ImageRateMultiplier *float64
BatchImageDiscountMultiplier *float64
BatchImageHoldMultiplier *float64
VideoRateIndependent bool
VideoRateMultiplier *float64
// 高峰时段倍率配置(PeakRateMultiplier 为 nil 时按 1.0 处理)
PeakRateEnabled bool
PeakStart string
PeakEnd string
PeakRateMultiplier *float64
ImagePrice1K *float64
ImagePrice2K *float64
ImagePrice4K *float64
VideoPrice480P *float64
VideoPrice720P *float64
VideoPrice1080P *float64
// VideoModelPrices 可选按模型族×分辨率覆盖视频每秒单价。
VideoModelPrices map[string]map[string]float64
// Codex alpha/search 网页搜索单次价格(USD/次,仅 openai 平台使用);nil/负数按默认价 0.01 处理
WebSearchPricePerCall *float64
// 搜索工具单价 per 1k
SearchPricePer1k *float64
// Grok Voice 显式定价(分组级)
AudioRealtimePricePerMin *float64
AudioTTSPricePerMillionChars *float64
AudioSTTPricePerHour *float64
ClaudeCodeOnly bool // 仅允许 Claude Code 客户端
FallbackGroupID *int64 // 降级分组 ID
// 无效请求兜底分组 ID(仅 anthropic 平台使用)
FallbackGroupIDOnInvalidRequest *int64
// 模型路由配置(仅 anthropic 平台使用)
ModelRouting map[string][]int64
ModelRoutingEnabled bool // 是否启用模型路由
MCPXMLInject *bool
// 支持的模型系列(仅 antigravity 平台使用)
SupportedModelScopes []string
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch bool
AllowLive bool
DefaultMappedModel string
RequireOAuthOnly bool
RequirePrivacySet bool
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig
ModelsListConfig GroupModelsListConfig
// RPMLimit 分组 RPM 上限(0 = 不限制)
RPMLimit int
// MaxReasoningEffort OpenAI/Codex 请求的推理强度上限,空字符串表示不限制。
MaxReasoningEffort string
// ReasoningEffortMappings OpenAI/Codex 推理强度精确映射。
ReasoningEffortMappings []ReasoningEffortMapping
// 分组利润控制(五个 token 平台分组可启用;margin/buffer 为小数,nil 按 0 处理)
ProfitControlEnabled bool
ProfitMinMargin *float64
ProfitSafetyBuffer *float64
// 从指定分组复制账号(创建分组后在同一事务内绑定)
CopyAccountsFromGroupIDs []int64
}
type UpdateGroupInput struct {
Name string
Description *string
Platform string
RateMultiplier *float64 // 使用指针以支持设置为0
IsExclusive *bool
Status string
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
LongContextPricingEnabled *bool
ModelPricing *[]ChannelModelPricing
// 图片生成计费配置(仅 antigravity 平台使用)
AllowImageGeneration *bool
AllowBatchImageGeneration *bool
ImageRateIndependent *bool
ImageRateMultiplier *float64
BatchImageDiscountMultiplier *float64
BatchImageHoldMultiplier *float64
VideoRateIndependent *bool
VideoRateMultiplier *float64
// 高峰时段倍率配置(nil 表示不修改)
PeakRateEnabled *bool
PeakStart *string
PeakEnd *string
PeakRateMultiplier *float64
ImagePrice1K *float64
ImagePrice2K *float64
ImagePrice4K *float64
VideoPrice480P *float64
VideoPrice720P *float64
VideoPrice1080P *float64
// VideoModelPrices 可选按模型族×分辨率覆盖;nil 表示不修改,空 map 表示清除。
VideoModelPrices map[string]map[string]float64
// Codex alpha/search 网页搜索单次价格(USD/次);nil 表示不修改,负数表示清除回默认价 0.01
WebSearchPricePerCall *float64
// 搜索工具单价;nil 不修改,负数清除
SearchPricePer1k *float64
// Grok Voice 显式定价;nil 表示不修改,负数表示清除
AudioRealtimePricePerMin *float64
AudioTTSPricePerMillionChars *float64
AudioSTTPricePerHour *float64
ClaudeCodeOnly *bool // 仅允许 Claude Code 客户端
FallbackGroupID *int64 // 降级分组 ID
// 无效请求兜底分组 ID(仅 anthropic 平台使用)
FallbackGroupIDOnInvalidRequest *int64
// 模型路由配置(仅 anthropic 平台使用)
ModelRouting map[string][]int64
ModelRoutingEnabled *bool // 是否启用模型路由
MCPXMLInject *bool
// 支持的模型系列(仅 antigravity 平台使用)
SupportedModelScopes *[]string
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch *bool
AllowLive *bool
DefaultMappedModel *string
RequireOAuthOnly *bool
RequirePrivacySet *bool
MessagesDispatchModelConfig *OpenAIMessagesDispatchModelConfig
ModelsListConfig *GroupModelsListConfig
// RPMLimit 分组 RPM 上限(0 = 不限制),nil 表示未提供不改动。
RPMLimit *int
// MaxReasoningEffort 空字符串表示清除上限;nil 表示未提供不改动。
MaxReasoningEffort *string
// ReasoningEffortMappings nil 表示不修改,空数组表示清空,非空数组表示替换。
ReasoningEffortMappings *[]ReasoningEffortMapping
// 分组利润控制(nil 表示不修改;margin/buffer 为小数)
ProfitControlEnabled *bool
ProfitMinMargin *float64
ProfitSafetyBuffer *float64
// 从指定分组复制账号(同步操作:先清空当前分组的账号绑定,再绑定源分组的账号)
CopyAccountsFromGroupIDs []int64
}
type CreateAccountInput struct {
Name string
Notes *string
Platform string
Type string
Credentials map[string]any
Extra map[string]any
ProxyID *int64
Concurrency int
Priority int
RateMultiplier *float64 // 账号计费倍率(>=0,允许 0)
LoadFactor *int
GroupIDs []int64
ExpiresAt *int64
AutoPauseOnExpired *bool
ProbeEnabled *bool
// SkipDefaultGroupBind prevents auto-binding to platform default group when GroupIDs is empty.
SkipDefaultGroupBind bool
// SkipMixedChannelCheck skips the mixed channel risk check when binding groups.
// This should only be set when the caller has explicitly confirmed the risk.
SkipMixedChannelCheck bool
}
// ShadowOptions is the input for CreateShadow.
// The shadow holds no credentials — the scheduler transparently delegates to the parent account's tokens.
type ShadowOptions struct {
Name string
Priority int
Concurrency int
GroupIDs []int64
}
type UpdateAccountInput struct {
Name string
Notes *string
Type string // Account type: oauth, setup-token, apikey
Credentials map[string]any
Extra map[string]any
ProxyID *int64
Concurrency *int // 使用指针区分"未提供"和"设置为0"
Priority *int // 使用指针区分"未提供"和"设置为0"
RateMultiplier *float64 // 账号计费倍率(>=0,允许 0)
LoadFactor *int
Status string
GroupIDs *[]int64
ExpiresAt *int64
AutoPauseOnExpired *bool
ProbeEnabled *bool
RateSyncEnabled *bool
SkipMixedChannelCheck bool // 跳过混合渠道检查(用户已确认风险)
}
// BulkUpdateAccountsInput describes the payload for bulk updating accounts.
type BulkUpdateAccountsInput struct {
AccountIDs []int64
Filters *BulkUpdateAccountFilters
Name string
ProxyID *int64
Concurrency *int
Priority *int
RateMultiplier *float64 // 账号计费倍率(>=0,允许 0)
LoadFactor *int
Status string
Schedulable *bool
GroupIDs *[]int64
Credentials map[string]any
Extra map[string]any
ProbeEnabled *bool
// SkipMixedChannelCheck skips the mixed channel risk check when binding groups.
// This should only be set when the caller has explicitly confirmed the risk.
SkipMixedChannelCheck bool
}
type BulkUpdateAccountFilters struct {
Platform string
Type string
Status string
Group string
Search string
PrivacyMode string
}
// BulkUpdateAccountResult captures the result for a single account update.
type BulkUpdateAccountResult struct {
AccountID int64 `json:"account_id"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
}
// AdminUpdateAPIKeyGroupIDResult is the result of AdminUpdateAPIKeyGroupID.
type AdminUpdateAPIKeyGroupIDResult struct {
APIKey *APIKey
AutoGrantedGroupAccess bool // true if a new exclusive group permission was auto-added
GrantedGroupID *int64 // the group ID that was auto-granted
GrantedGroupName string // the group name that was auto-granted
}
// ReplaceUserGroupResult 分组替换操作的结果
type ReplaceUserGroupResult struct {
MigratedKeys int64 // 迁移的 Key 数量
}
// UserRPMStatus describes a user's current per-minute RPM usage.
type UserRPMStatus struct {
UserRPMUsed int `json:"user_rpm_used"`
UserRPMLimit int `json:"user_rpm_limit"`
PerGroup []UserGroupRPMStatus `json:"per_group"`
}
// UserGroupRPMStatus describes current per-minute RPM usage for one user/group pair.
type UserGroupRPMStatus struct {
GroupID int64 `json:"group_id"`
GroupName string `json:"group_name"`
Used int `json:"used"`
Limit int `json:"limit"`
Source string `json:"source"` // "group" | "override"
}
// BulkUpdateAccountsResult is the aggregated response for bulk updates.
type BulkUpdateAccountsResult struct {
Success int `json:"success"`
Failed int `json:"failed"`
SuccessIDs []int64 `json:"success_ids"`
FailedIDs []int64 `json:"failed_ids"`
Results []BulkUpdateAccountResult `json:"results"`
LongContextInheritedCount int `json:"long_context_inherited_count,omitempty"`
}
type CreateProxyInput struct {
Name string
Protocol string
Host string
Port int
Username string
Password string
ExpiresAt *time.Time
FallbackMode string
BackupProxyID *int64
ExpiryWarnDays int
}
type UpdateProxyInput struct {
Name string
Protocol string
Host string
Port int
Username string
Password string
Status string
ExpiresAt *time.Time
FallbackMode string
BackupProxyID *int64
ExpiryWarnDays int
}
type GenerateRedeemCodesInput struct {
Count int
Type string
Value float64
GroupID *int64 // 订阅类型专用:关联的分组ID
ValidityDays int // 订阅类型专用:有效天数
ExpiresAt *time.Time
}
type ProxyBatchDeleteResult struct {
DeletedIDs []int64 `json:"deleted_ids"`
Skipped []ProxyBatchDeleteSkipped `json:"skipped"`
}
type ProxyBatchDeleteSkipped struct {
ID int64 `json:"id"`
Reason string `json:"reason"`
}
// ProxyTestResult represents the result of testing a proxy
type ProxyTestResult struct {
Success bool `json:"success"`
Message string `json:"message"`
LatencyMs int64 `json:"latency_ms,omitempty"`
IPAddress string `json:"ip_address,omitempty"`
City string `json:"city,omitempty"`
Region string `json:"region,omitempty"`
Country string `json:"country,omitempty"`
CountryCode string `json:"country_code,omitempty"`
}
type ProxyQualityCheckResult struct {
ProxyID int64 `json:"proxy_id"`
Score int `json:"score"`
Grade string `json:"grade"`
Summary string `json:"summary"`
ExitIP string `json:"exit_ip,omitempty"`
Country string `json:"country,omitempty"`
CountryCode string `json:"country_code,omitempty"`
BaseLatencyMs int64 `json:"base_latency_ms,omitempty"`
PassedCount int `json:"passed_count"`
WarnCount int `json:"warn_count"`
FailedCount int `json:"failed_count"`
ChallengeCount int `json:"challenge_count"`
CheckedAt int64 `json:"checked_at"`
Items []ProxyQualityCheckItem `json:"items"`
}
type ProxyQualityCheckItem struct {
Target string `json:"target"`
Status string `json:"status"` // pass/warn/fail/challenge
HTTPStatus int `json:"http_status,omitempty"`
LatencyMs int64 `json:"latency_ms,omitempty"`
Message string `json:"message,omitempty"`
CFRay string `json:"cf_ray,omitempty"`
}
// ProxyExitInfo represents proxy exit information from ip-api.com
type ProxyExitInfo struct {
IP string
City string
Region string
Country string
CountryCode string
}
// ProxyExitInfoProber tests proxy connectivity and retrieves exit information
type ProxyExitInfoProber interface {
ProbeProxy(ctx context.Context, proxyURL string) (*ProxyExitInfo, int64, error)
}
type groupExistenceBatchReader interface {
ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error)
}
type proxyQualityTarget struct {
Target string
URL string
Method string
AllowedStatuses map[int]struct{}
}
var proxyQualityTargets = []proxyQualityTarget{
{
Target: "openai",
URL: "https://api.openai.com/v1/models",
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
},
},
{
Target: "anthropic",
URL: "https://api.anthropic.com/v1/messages",
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
http.StatusMethodNotAllowed: {},
http.StatusNotFound: {},
http.StatusBadRequest: {},
},
},
{
Target: "gemini",
URL: "https://generativelanguage.googleapis.com/$discovery/rest?version=v1beta",
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusOK: {},
},
},
{
Target: "grok",
URL: "https://api.x.ai/v1/models",
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
},
},
}
const (
proxyQualityRequestTimeout = 15 * time.Second
proxyQualityResponseHeaderTimeout = 10 * time.Second
proxyQualityMaxBodyBytes = int64(8 * 1024)
proxyQualityClientUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36"
)
var ErrRPMStatusUnavailable = infraerrors.New(http.StatusNotImplemented, "RPM_STATUS_UNAVAILABLE", "RPM cache not available")
// adminServiceImpl implements AdminService
type adminServiceImpl struct {
userRepo UserRepository
groupRepo GroupRepository
groupDuplicateRepo GroupDuplicateRepository
accountRepo AccountRepository
accountDuplicateRepo AccountDuplicateRepository
accountBillingRepo AccountBillingSettingsRepository
proxyRepo ProxyRepository
apiKeyRepo APIKeyRepository
redeemCodeRepo RedeemCodeRepository
userGroupRateRepo UserGroupRateRepository
userRPMCache UserRPMCache
billingCacheService *BillingCacheService
proxyProber ProxyExitInfoProber
proxyLatencyCache ProxyLatencyCache
authCacheInvalidator APIKeyAuthCacheInvalidator
entClient *dbent.Client // 用于开启数据库事务
settingService *SettingService
defaultSubAssigner DefaultSubscriptionAssigner
userSubRepo UserSubscriptionRepository
privacyClientFactory PrivacyClientFactory
runtimeBlocker AccountRuntimeBlocker
affiliateService adminRechargeAffiliateAccruer
compositeRouteRepo CompositeModelRouteRepository
compositeResolver *CompositeRouteResolver
// 分组平台变更后用来失效渠道缓存;可为 nil(缓存会在 TTL 到期后自然重建)
channelCacheInvalidator ChannelCacheInvalidator
}
// ChannelCacheInvalidator 失效渠道缓存。
// 窄接口,避免 admin 服务依赖整个 ChannelService——与 APIKeyAuthCacheInvalidator 同一思路。
type ChannelCacheInvalidator interface {
InvalidateCache()
}
type adminRechargeAffiliateAccruer interface {
AccrueInviteRebate(ctx context.Context, inviteeUserID int64, baseRechargeAmount float64) (float64, error)
}
type userGroupRateBatchReader interface {
GetByUserIDs(ctx context.Context, userIDs []int64) (map[int64]map[int64]float64, error)
}
// NewAdminService creates a new AdminService
func NewAdminService(
userRepo UserRepository,
groupRepo AdminGroupRepository,
accountRepo AdminAccountRepository,
proxyRepo ProxyRepository,
apiKeyRepo APIKeyRepository,
redeemCodeRepo RedeemCodeRepository,
userGroupRateRepo UserGroupRateRepository,
userRPMCache UserRPMCache,
billingCacheService *BillingCacheService,
proxyProber ProxyExitInfoProber,
proxyLatencyCache ProxyLatencyCache,
authCacheInvalidator APIKeyAuthCacheInvalidator,
entClient *dbent.Client,
settingService *SettingService,
defaultSubAssigner DefaultSubscriptionAssigner,
userSubRepo UserSubscriptionRepository,
privacyClientFactory PrivacyClientFactory,
runtimeBlocker AccountRuntimeBlocker,
affiliateService *AffiliateService,
compositeRouteRepo CompositeModelRouteRepository,
compositeResolver *CompositeRouteResolver,
channelCacheInvalidator ChannelCacheInvalidator,
) AdminService {
return &adminServiceImpl{
userRepo: userRepo,
groupRepo: groupRepo,
groupDuplicateRepo: groupRepo,
accountRepo: accountRepo,
accountDuplicateRepo: accountRepo,
accountBillingRepo: accountRepo,
proxyRepo: proxyRepo,
apiKeyRepo: apiKeyRepo,
redeemCodeRepo: redeemCodeRepo,
userGroupRateRepo: userGroupRateRepo,
userRPMCache: userRPMCache,
billingCacheService: billingCacheService,
proxyProber: proxyProber,
proxyLatencyCache: proxyLatencyCache,
authCacheInvalidator: authCacheInvalidator,
entClient: entClient,
settingService: settingService,
defaultSubAssigner: defaultSubAssigner,
userSubRepo: userSubRepo,
privacyClientFactory: privacyClientFactory,
runtimeBlocker: runtimeBlocker,
affiliateService: affiliateService,
compositeRouteRepo: compositeRouteRepo,
compositeResolver: compositeResolver,
channelCacheInvalidator: channelCacheInvalidator,
}
}
@@ -0,0 +1,567 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// Stubs
// ---------------------------------------------------------------------------
// userRepoStubForGroupUpdate implements UserRepository for AdminUpdateAPIKeyGroupID tests.
type userRepoStubForGroupUpdate struct {
addGroupErr error
addGroupCalled bool
addedUserID int64
addedGroupID int64
}
func (s *userRepoStubForGroupUpdate) AddGroupToAllowedGroups(_ context.Context, userID int64, groupID int64) error {
s.addGroupCalled = true
s.addedUserID = userID
s.addedGroupID = groupID
return s.addGroupErr
}
func (s *userRepoStubForGroupUpdate) Create(context.Context, *User) error { panic("unexpected") }
func (s *userRepoStubForGroupUpdate) CreateWithEmailAliasGuard(context.Context, *User) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetByID(context.Context, int64) (*User, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetByEmail(context.Context, string) (*User, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetFirstAdmin(context.Context) (*User, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) Update(context.Context, *User, UserUpdateFields) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) Delete(context.Context, int64) error { panic("unexpected") }
func (s *userRepoStubForGroupUpdate) GetUserAvatar(context.Context, int64) (*UserAvatar, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) UpsertUserAvatar(context.Context, int64, UpsertUserAvatarInput) (*UserAvatar, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) DeleteUserAvatar(context.Context, int64) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) List(context.Context, pagination.PaginationParams) ([]User, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) ListWithFilters(context.Context, pagination.PaginationParams, UserListFilters) ([]User, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) UpdateBalance(context.Context, int64, float64) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) DeductBalance(context.Context, int64, float64) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) AdjustBalance(ctx context.Context, id int64, delta float64) (BalanceChange, error) {
panic("unexpected AdjustBalance call")
}
func (s *userRepoStubForGroupUpdate) SetBalance(ctx context.Context, id int64, value float64) (BalanceChange, error) {
panic("unexpected SetBalance call")
}
func (s *userRepoStubForGroupUpdate) UpdateConcurrency(context.Context, int64, int) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) BatchSetConcurrency(context.Context, []int64, int) (int, error) {
return 0, nil
}
func (s *userRepoStubForGroupUpdate) BatchAddConcurrency(context.Context, []int64, int) (int, error) {
return 0, nil
}
func (s *userRepoStubForGroupUpdate) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
return 0, nil
}
func (s *userRepoStubForGroupUpdate) ExistsByEmail(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) ExistsByEmailAlias(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) UpdateTotpSecret(context.Context, int64, *string) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) EnableTotp(context.Context, int64) error { panic("unexpected") }
func (s *userRepoStubForGroupUpdate) DisableTotp(context.Context, int64) error { panic("unexpected") }
func (s *userRepoStubForGroupUpdate) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) {
panic("unexpected GetByIDIncludeDeleted call")
}
func (s *userRepoStubForGroupUpdate) ListUserAuthIdentities(context.Context, int64) ([]UserAuthIdentityRecord, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) UnbindUserAuthProvider(context.Context, int64, string) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetLatestUsedAtByUserIDs(context.Context, []int64) (map[int64]*time.Time, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetLatestUsedAtByUserID(context.Context, int64) (*time.Time, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) UpdateUserLastActiveAt(context.Context, int64, time.Time) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) RemoveGroupFromUserAllowedGroups(context.Context, int64, int64) error {
panic("unexpected")
}
// apiKeyRepoStubForGroupUpdate implements APIKeyRepository for AdminUpdateAPIKeyGroupID tests.
type apiKeyRepoStubForGroupUpdate struct {
key *APIKey
getErr error
updateErr error
updated *APIKey // captures what was passed to Update
}
func (s *apiKeyRepoStubForGroupUpdate) GetByID(_ context.Context, _ int64) (*APIKey, error) {
if s.getErr != nil {
return nil, s.getErr
}
clone := *s.key
return &clone, nil
}
func (s *apiKeyRepoStubForGroupUpdate) Update(_ context.Context, key *APIKey, _ APIKeyUpdateFields) error {
if s.updateErr != nil {
return s.updateErr
}
clone := *key
s.updated = &clone
return nil
}
// Unused methods panic on unexpected call.
func (s *apiKeyRepoStubForGroupUpdate) Create(context.Context, *APIKey) error { panic("unexpected") }
func (s *apiKeyRepoStubForGroupUpdate) GetKeyAndOwnerID(context.Context, int64) (string, int64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) GetByKey(context.Context, string) (*APIKey, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) GetByKeyForAuth(context.Context, string) (*APIKey, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) Delete(context.Context, int64) error { panic("unexpected") }
func (s *apiKeyRepoStubForGroupUpdate) DeleteWithAudit(context.Context, int64) error {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ListByUserID(context.Context, int64, pagination.PaginationParams, APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) VerifyOwnership(context.Context, int64, []int64) ([]int64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) CountByUserID(context.Context, int64) (int64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ExistsByKey(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ListByGroupID(context.Context, int64, pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) SearchAPIKeys(context.Context, int64, string, int) ([]APIKey, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ClearGroupIDByGroupID(context.Context, int64) (int64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) CountByGroupID(context.Context, int64) (int64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ListKeysByUserID(context.Context, int64) ([]string, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ListKeysByGroupID(context.Context, int64) ([]string, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) IncrementQuotaUsed(context.Context, int64, float64) (float64, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) UpdateLastUsed(context.Context, int64, time.Time) error {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) IncrementRateLimitUsage(context.Context, int64, float64) error {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) ResetRateLimitWindows(context.Context, int64) error {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) {
panic("unexpected")
}
func (s *apiKeyRepoStubForGroupUpdate) UpdateGroupIDByUserAndGroup(context.Context, int64, int64, int64) (int64, error) {
panic("unexpected")
}
// groupRepoStubForGroupUpdate implements GroupRepository for AdminUpdateAPIKeyGroupID tests.
type groupRepoStubForGroupUpdate struct {
group *Group
getErr error
lastGetByIDArg int64
}
func (s *groupRepoStubForGroupUpdate) GetByID(_ context.Context, id int64) (*Group, error) {
s.lastGetByIDArg = id
if s.getErr != nil {
return nil, s.getErr
}
clone := *s.group
return &clone, nil
}
// Unused methods panic on unexpected call.
func (s *groupRepoStubForGroupUpdate) Create(context.Context, *Group) error { panic("unexpected") }
func (s *groupRepoStubForGroupUpdate) GetByIDLite(context.Context, int64) (*Group, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) Update(context.Context, *Group) error { panic("unexpected") }
func (s *groupRepoStubForGroupUpdate) Delete(context.Context, int64) error { panic("unexpected") }
func (s *groupRepoStubForGroupUpdate) DeleteCascade(context.Context, int64) ([]int64, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) List(context.Context, pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string, *bool) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) ListActive(context.Context) ([]Group, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) ListActiveByPlatform(context.Context, string) ([]Group, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) ExistsByName(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) GetAccountCount(context.Context, int64) (int64, int64, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) DeleteAccountGroupsByGroupID(context.Context, int64) (int64, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) GetAccountIDsByGroupIDs(context.Context, []int64) ([]int64, error) {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) BindAccountsToGroup(context.Context, int64, []int64) error {
panic("unexpected")
}
func (s *groupRepoStubForGroupUpdate) UpdateSortOrders(context.Context, []GroupSortOrderUpdate) error {
panic("unexpected")
}
type userSubRepoStubForGroupUpdate struct {
userSubRepoNoop
getActiveSub *UserSubscription
getActiveErr error
called bool
calledUserID int64
calledGroupID int64
}
func (s *userSubRepoStubForGroupUpdate) GetActiveByUserIDAndGroupID(_ context.Context, userID, groupID int64) (*UserSubscription, error) {
s.called = true
s.calledUserID = userID
s.calledGroupID = groupID
if s.getActiveErr != nil {
return nil, s.getActiveErr
}
if s.getActiveSub == nil {
return nil, ErrSubscriptionNotFound
}
clone := *s.getActiveSub
return &clone, nil
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
func TestAdminService_AdminUpdateAPIKeyGroupID_KeyNotFound(t *testing.T) {
repo := &apiKeyRepoStubForGroupUpdate{getErr: ErrAPIKeyNotFound}
svc := &adminServiceImpl{apiKeyRepo: repo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 999, int64Ptr(1))
require.ErrorIs(t, err, ErrAPIKeyNotFound)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_NilGroupID_NoOp(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: int64Ptr(5)}
repo := &apiKeyRepoStubForGroupUpdate{key: existing}
svc := &adminServiceImpl{apiKeyRepo: repo}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, nil)
require.NoError(t, err)
require.Equal(t, int64(1), got.APIKey.ID)
// Update should NOT have been called (updated stays nil)
require.Nil(t, repo.updated)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_Unbind(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: int64Ptr(5), Group: &Group{ID: 5, Name: "Old"}}
repo := &apiKeyRepoStubForGroupUpdate{key: existing}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: repo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(0))
require.NoError(t, err)
require.Nil(t, got.APIKey.GroupID, "group_id should be nil after unbind")
require.Nil(t, got.APIKey.Group, "group object should be nil after unbind")
require.NotNil(t, repo.updated, "Update should have been called")
require.Nil(t, repo.updated.GroupID)
require.Equal(t, []string{"sk-test"}, cache.keys, "cache should be invalidated")
}
func TestAdminService_AdminUpdateAPIKeyGroupID_BindActiveGroup(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Pro", Status: StatusActive}}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
require.Equal(t, int64(10), *got.APIKey.GroupID)
require.Equal(t, int64(10), *apiKeyRepo.updated.GroupID)
require.Equal(t, []string{"sk-test"}, cache.keys)
// M3: verify correct group ID was passed to repo
require.Equal(t, int64(10), groupRepo.lastGetByIDArg)
// C1 fix: verify Group object is populated
require.NotNil(t, got.APIKey.Group)
require.Equal(t, "Pro", got.APIKey.Group.Name)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_SameGroup_Idempotent(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: int64Ptr(10), Group: &Group{ID: 10, Name: "Pro"}}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Pro", Status: StatusActive}}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
require.Equal(t, int64(10), *got.APIKey.GroupID)
// Update is still called (current impl doesn't short-circuit on same group)
require.NotNil(t, apiKeyRepo.updated)
require.Equal(t, []string{"sk-test"}, cache.keys)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_GroupNotFound(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test"}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{getErr: ErrGroupNotFound}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(99))
require.ErrorIs(t, err, ErrGroupNotFound)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_GroupNotActive(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test"}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 5, Status: StatusDisabled}}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(5))
require.Error(t, err)
require.Equal(t, "GROUP_NOT_ACTIVE", infraerrors.Reason(err))
}
func TestAdminService_AdminUpdateAPIKeyGroupID_UpdateFails(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: int64Ptr(3)}
repo := &apiKeyRepoStubForGroupUpdate{key: existing, updateErr: errors.New("db write error")}
svc := &adminServiceImpl{apiKeyRepo: repo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(0))
require.Error(t, err)
require.Contains(t, err.Error(), "update api key")
}
func TestAdminService_AdminUpdateAPIKeyGroupID_NegativeGroupID(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test"}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(-5))
require.Error(t, err)
require.Equal(t, "INVALID_GROUP_ID", infraerrors.Reason(err))
}
func TestAdminService_AdminUpdateAPIKeyGroupID_PointerIsolation(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Pro", Status: StatusActive}}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, authCacheInvalidator: cache}
inputGID := int64(10)
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, &inputGID)
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
// Mutating the input pointer must NOT affect the stored value
inputGID = 999
require.Equal(t, int64(10), *got.APIKey.GroupID)
require.Equal(t, int64(10), *apiKeyRepo.updated.GroupID)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_NilCacheInvalidator(t *testing.T) {
existing := &APIKey{ID: 1, Key: "sk-test"}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 7, Status: StatusActive}}
// authCacheInvalidator is nil should not panic
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(7))
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
require.Equal(t, int64(7), *got.APIKey.GroupID)
}
// ---------------------------------------------------------------------------
// Tests: AllowedGroup auto-sync
// ---------------------------------------------------------------------------
func TestAdminService_AdminUpdateAPIKeyGroupID_ExclusiveGroup_AddsAllowedGroup(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Exclusive", Status: StatusActive, IsExclusive: true, SubscriptionType: SubscriptionTypeStandard}}
userRepo := &userRepoStubForGroupUpdate{}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
require.Equal(t, int64(10), *got.APIKey.GroupID)
// 验证 AddGroupToAllowedGroups 被调用,且参数正确
require.True(t, userRepo.addGroupCalled)
require.Equal(t, int64(42), userRepo.addedUserID)
require.Equal(t, int64(10), userRepo.addedGroupID)
// 验证 result 标记了自动授权
require.True(t, got.AutoGrantedGroupAccess)
require.NotNil(t, got.GrantedGroupID)
require.Equal(t, int64(10), *got.GrantedGroupID)
require.Equal(t, "Exclusive", got.GrantedGroupName)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_NonExclusiveGroup_NoAllowedGroupUpdate(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Public", Status: StatusActive, IsExclusive: false, SubscriptionType: SubscriptionTypeStandard}}
userRepo := &userRepoStubForGroupUpdate{}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.NoError(t, err)
require.NotNil(t, got.APIKey.GroupID)
// 非专属分组不触发 AddGroupToAllowedGroups
require.False(t, userRepo.addGroupCalled)
require.False(t, got.AutoGrantedGroupAccess)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_SubscriptionGroup_Blocked(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Sub", Status: StatusActive, IsExclusive: false, SubscriptionType: SubscriptionTypeSubscription}}
userRepo := &userRepoStubForGroupUpdate{}
userSubRepo := &userSubRepoStubForGroupUpdate{getActiveErr: ErrSubscriptionNotFound}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo, userSubRepo: userSubRepo}
// 无有效订阅时应拒绝绑定
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.Error(t, err)
require.Equal(t, "SUBSCRIPTION_REQUIRED", infraerrors.Reason(err))
require.True(t, userSubRepo.called)
require.Equal(t, int64(42), userSubRepo.calledUserID)
require.Equal(t, int64(10), userSubRepo.calledGroupID)
require.False(t, userRepo.addGroupCalled)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_SubscriptionGroup_RequiresRepo(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Sub", Status: StatusActive, IsExclusive: false, SubscriptionType: SubscriptionTypeSubscription}}
userRepo := &userRepoStubForGroupUpdate{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo}
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.Error(t, err)
require.Equal(t, "SUBSCRIPTION_REPOSITORY_UNAVAILABLE", infraerrors.Reason(err))
require.False(t, userRepo.addGroupCalled)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_SubscriptionGroup_AllowsActiveSubscription(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Sub", Status: StatusActive, IsExclusive: true, SubscriptionType: SubscriptionTypeSubscription}}
userRepo := &userRepoStubForGroupUpdate{}
userSubRepo := &userSubRepoStubForGroupUpdate{
getActiveSub: &UserSubscription{ID: 99, UserID: 42, GroupID: 10},
}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo, userSubRepo: userSubRepo}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.NoError(t, err)
require.True(t, userSubRepo.called)
require.NotNil(t, got.APIKey.GroupID)
require.Equal(t, int64(10), *got.APIKey.GroupID)
require.False(t, userRepo.addGroupCalled)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_ExclusiveGroup_AllowedGroupAddFails_ReturnsError(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: nil}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
groupRepo := &groupRepoStubForGroupUpdate{group: &Group{ID: 10, Name: "Exclusive", Status: StatusActive, IsExclusive: true, SubscriptionType: SubscriptionTypeStandard}}
userRepo := &userRepoStubForGroupUpdate{addGroupErr: errors.New("db error")}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, groupRepo: groupRepo, userRepo: userRepo}
// 严格模式:AddGroupToAllowedGroups 失败时,整体操作报错
_, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(10))
require.Error(t, err)
require.Contains(t, err.Error(), "add group to user allowed groups")
require.True(t, userRepo.addGroupCalled)
// apiKey 不应被更新
require.Nil(t, apiKeyRepo.updated)
}
func TestAdminService_AdminUpdateAPIKeyGroupID_Unbind_NoAllowedGroupUpdate(t *testing.T) {
existing := &APIKey{ID: 1, UserID: 42, Key: "sk-test", GroupID: int64Ptr(10), Group: &Group{ID: 10, Name: "Exclusive"}}
apiKeyRepo := &apiKeyRepoStubForGroupUpdate{key: existing}
userRepo := &userRepoStubForGroupUpdate{}
cache := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{apiKeyRepo: apiKeyRepo, userRepo: userRepo, authCacheInvalidator: cache}
got, err := svc.AdminUpdateAPIKeyGroupID(context.Background(), 1, int64Ptr(0))
require.NoError(t, err)
require.Nil(t, got.APIKey.GroupID)
// 解绑时不修改 allowed_groups
require.False(t, userRepo.addGroupCalled)
require.False(t, got.AutoGrantedGroupAccess)
}
@@ -0,0 +1,302 @@
//go:build unit
package service
import (
"context"
"database/sql"
"testing"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/authidentity"
"github.com/Wei-Shaw/sub2api/ent/authidentitychannel"
"github.com/Wei-Shaw/sub2api/ent/enttest"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
_ "modernc.org/sqlite"
)
func newAdminServiceAuthIdentityBindingTestClient(t *testing.T) *dbent.Client {
t.Helper()
db, err := sql.Open("sqlite", "file:admin_service_auth_identity_binding?mode=memory&cache=shared&_fk=1")
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
_, err = db.Exec("PRAGMA foreign_keys = ON")
require.NoError(t, err)
drv := entsql.OpenDB(dialect.SQLite, db)
client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(drv)))
t.Cleanup(func() { _ = client.Close() })
return client
}
func TestAdminServiceBindUserAuthIdentityCreatesCanonicalAndChannelBinding(t *testing.T) {
client := newAdminServiceAuthIdentityBindingTestClient(t)
ctx := context.Background()
user, err := client.User.Create().
SetEmail("bind-target@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
svc := &adminServiceImpl{
userRepo: &userRepoStub{user: &User{ID: user.ID, Email: user.Email, Status: StatusActive}},
entClient: client,
}
result, err := svc.BindUserAuthIdentity(ctx, user.ID, AdminBindAuthIdentityInput{
ProviderType: "wechat",
ProviderKey: "wechat-main",
ProviderSubject: "union-123",
Metadata: map[string]any{"source": "admin-repair"},
Channel: &AdminBindAuthIdentityChannelInput{
Channel: "open",
ChannelAppID: "wx-open",
ChannelSubject: "openid-123",
Metadata: map[string]any{"scene": "migration"},
},
})
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, user.ID, result.UserID)
require.Equal(t, "wechat", result.ProviderType)
require.Equal(t, "wechat-main", result.ProviderKey)
require.NotNil(t, result.VerifiedAt)
require.NotNil(t, result.Channel)
require.Equal(t, "open", result.Channel.Channel)
identity, err := client.AuthIdentity.Query().
Where(
authidentity.ProviderTypeEQ("wechat"),
authidentity.ProviderKeyEQ("wechat-main"),
authidentity.ProviderSubjectEQ("union-123"),
).
Only(ctx)
require.NoError(t, err)
require.Equal(t, user.ID, identity.UserID)
require.Equal(t, "admin-repair", identity.Metadata["source"])
require.NotNil(t, identity.VerifiedAt)
channel, err := client.AuthIdentityChannel.Query().
Where(
authidentitychannel.ProviderTypeEQ("wechat"),
authidentitychannel.ProviderKeyEQ("wechat-main"),
authidentitychannel.ChannelEQ("open"),
authidentitychannel.ChannelAppIDEQ("wx-open"),
authidentitychannel.ChannelSubjectEQ("openid-123"),
).
Only(ctx)
require.NoError(t, err)
require.Equal(t, identity.ID, channel.IdentityID)
require.Equal(t, "migration", channel.Metadata["scene"])
}
func TestAdminServiceBindUserAuthIdentityRejectsOtherOwner(t *testing.T) {
client := newAdminServiceAuthIdentityBindingTestClient(t)
ctx := context.Background()
owner, err := client.User.Create().
SetEmail("owner@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
target, err := client.User.Create().
SetEmail("target@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
_, err = client.AuthIdentity.Create().
SetUserID(owner.ID).
SetProviderType("oidc").
SetProviderKey("https://issuer.example").
SetProviderSubject("subject-1").
Save(ctx)
require.NoError(t, err)
svc := &adminServiceImpl{
userRepo: &userRepoStub{user: &User{ID: target.ID, Email: target.Email, Status: StatusActive}},
entClient: client,
}
_, err = svc.BindUserAuthIdentity(ctx, target.ID, AdminBindAuthIdentityInput{
ProviderType: "oidc",
ProviderKey: "https://issuer.example",
ProviderSubject: "subject-1",
})
require.Error(t, err)
require.Equal(t, "AUTH_IDENTITY_OWNERSHIP_CONFLICT", infraerrors.Reason(err))
}
func TestAdminServiceBindUserAuthIdentityIsIdempotentForSameUser(t *testing.T) {
client := newAdminServiceAuthIdentityBindingTestClient(t)
ctx := context.Background()
user, err := client.User.Create().
SetEmail("same-user@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
svc := &adminServiceImpl{
userRepo: &userRepoStub{user: &User{ID: user.ID, Email: user.Email, Status: StatusActive}},
entClient: client,
}
first, err := svc.BindUserAuthIdentity(ctx, user.ID, AdminBindAuthIdentityInput{
ProviderType: "oidc",
ProviderKey: "https://issuer.example",
ProviderSubject: "subject-2",
Metadata: map[string]any{"source": "first"},
})
require.NoError(t, err)
second, err := svc.BindUserAuthIdentity(ctx, user.ID, AdminBindAuthIdentityInput{
ProviderType: "oidc",
ProviderKey: "https://issuer.example",
ProviderSubject: "subject-2",
Metadata: map[string]any{"source": "second"},
})
require.NoError(t, err)
require.Equal(t, first.UserID, second.UserID)
require.Equal(t, "second", second.Metadata["source"])
identities, err := client.AuthIdentity.Query().
Where(
authidentity.ProviderTypeEQ("oidc"),
authidentity.ProviderKeyEQ("https://issuer.example"),
authidentity.ProviderSubjectEQ("subject-2"),
).
All(ctx)
require.NoError(t, err)
require.Len(t, identities, 1)
require.Equal(t, "second", identities[0].Metadata["source"])
}
func TestAdminServiceBindUserAuthIdentityReusesLegacyWeChatAliasRecords(t *testing.T) {
client := newAdminServiceAuthIdentityBindingTestClient(t)
ctx := context.Background()
user, err := client.User.Create().
SetEmail("wechat-alias@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
legacyIdentity, err := client.AuthIdentity.Create().
SetUserID(user.ID).
SetProviderType("wechat").
SetProviderKey("wechat").
SetProviderSubject("union-legacy-123").
SetMetadata(map[string]any{"source": "legacy"}).
Save(ctx)
require.NoError(t, err)
legacyChannel, err := client.AuthIdentityChannel.Create().
SetIdentityID(legacyIdentity.ID).
SetProviderType("wechat").
SetProviderKey("wechat").
SetChannel("open").
SetChannelAppID("wx-open").
SetChannelSubject("openid-legacy-123").
SetMetadata(map[string]any{"scene": "legacy"}).
Save(ctx)
require.NoError(t, err)
svc := &adminServiceImpl{
userRepo: &userRepoStub{user: &User{ID: user.ID, Email: user.Email, Status: StatusActive}},
entClient: client,
}
result, err := svc.BindUserAuthIdentity(ctx, user.ID, AdminBindAuthIdentityInput{
ProviderType: "wechat",
ProviderKey: "wechat-main",
ProviderSubject: "union-legacy-123",
Metadata: map[string]any{"source": "admin-repair"},
Channel: &AdminBindAuthIdentityChannelInput{
Channel: "open",
ChannelAppID: "wx-open",
ChannelSubject: "openid-legacy-123",
Metadata: map[string]any{"scene": "admin-repair"},
},
})
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "wechat-main", result.ProviderKey)
require.NotNil(t, result.Channel)
require.Equal(t, "open", result.Channel.Channel)
identity, err := client.AuthIdentity.Get(ctx, legacyIdentity.ID)
require.NoError(t, err)
require.Equal(t, "wechat-main", identity.ProviderKey)
require.Equal(t, "admin-repair", identity.Metadata["source"])
channel, err := client.AuthIdentityChannel.Get(ctx, legacyChannel.ID)
require.NoError(t, err)
require.Equal(t, "wechat-main", channel.ProviderKey)
require.Equal(t, legacyIdentity.ID, channel.IdentityID)
require.Equal(t, "admin-repair", channel.Metadata["scene"])
identityCount, err := client.AuthIdentity.Query().
Where(
authidentity.ProviderTypeEQ("wechat"),
authidentity.ProviderSubjectEQ("union-legacy-123"),
).
Count(ctx)
require.NoError(t, err)
require.Equal(t, 1, identityCount)
channelCount, err := client.AuthIdentityChannel.Query().
Where(
authidentitychannel.ProviderTypeEQ("wechat"),
authidentitychannel.ChannelEQ("open"),
authidentitychannel.ChannelAppIDEQ("wx-open"),
authidentitychannel.ChannelSubjectEQ("openid-legacy-123"),
).
Count(ctx)
require.NoError(t, err)
require.Equal(t, 1, channelCount)
}
func TestAdminServiceBindUserAuthIdentityRejectsInvalidProviderType(t *testing.T) {
client := newAdminServiceAuthIdentityBindingTestClient(t)
ctx := context.Background()
user, err := client.User.Create().
SetEmail("invalid-provider@example.com").
SetPasswordHash("hash").
SetRole(RoleUser).
SetStatus(StatusActive).
Save(ctx)
require.NoError(t, err)
svc := &adminServiceImpl{
userRepo: &userRepoStub{user: &User{ID: user.ID, Email: user.Email, Status: StatusActive}},
entClient: client,
}
_, err = svc.BindUserAuthIdentity(ctx, user.ID, AdminBindAuthIdentityInput{
ProviderType: "github",
ProviderKey: "github-main",
ProviderSubject: "subject-3",
})
require.Error(t, err)
require.Equal(t, "INVALID_INPUT", infraerrors.Reason(err))
}
@@ -0,0 +1,92 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/require"
)
type batchLimitsUserRepoStub struct {
*userRepoStub
calls int
userIDs []int64
concurrency *int
rpmLimit *int
affected int
err error
}
func (s *batchLimitsUserRepoStub) BatchUpdateLimits(_ context.Context, userIDs []int64, concurrency, rpmLimit *int) (int, error) {
s.calls++
s.userIDs = append([]int64(nil), userIDs...)
s.concurrency = cloneBatchLimitValue(concurrency)
s.rpmLimit = cloneBatchLimitValue(rpmLimit)
return s.affected, s.err
}
func cloneBatchLimitValue(value *int) *int {
if value == nil {
return nil
}
cloned := *value
return &cloned
}
func TestAdminServiceBatchUpdateLimitsPassesOnlyProvidedFields(t *testing.T) {
concurrency := 0
repo := &batchLimitsUserRepoStub{
userRepoStub: &userRepoStub{},
affected: 2,
}
invalidator := &authCacheInvalidatorStub{}
service := &adminServiceImpl{userRepo: repo, authCacheInvalidator: invalidator}
affected, err := service.BatchUpdateLimits(
context.Background(),
[]int64{3, 0, 3, 7, -1},
&concurrency,
nil,
)
require.NoError(t, err)
require.Equal(t, 2, affected)
require.Equal(t, []int64{3, 7}, repo.userIDs)
require.Equal(t, pointerToInt(0), repo.concurrency)
require.Nil(t, repo.rpmLimit)
require.Equal(t, []int64{3, 7}, invalidator.userIDs)
}
func TestAdminServiceBatchUpdateLimitsDoesNotInvalidateCacheOnRepositoryError(t *testing.T) {
rpmLimit := 60
repo := &batchLimitsUserRepoStub{
userRepoStub: &userRepoStub{},
err: errors.New("database unavailable"),
}
invalidator := &authCacheInvalidatorStub{}
service := &adminServiceImpl{userRepo: repo, authCacheInvalidator: invalidator}
affected, err := service.BatchUpdateLimits(context.Background(), []int64{1, 2}, nil, &rpmLimit)
require.EqualError(t, err, "database unavailable")
require.Zero(t, affected)
require.Empty(t, invalidator.userIDs)
}
func TestAdminServiceBatchUpdateLimitsRequiresAField(t *testing.T) {
repo := &batchLimitsUserRepoStub{userRepoStub: &userRepoStub{}}
service := &adminServiceImpl{userRepo: repo, authCacheInvalidator: &authCacheInvalidatorStub{}}
affected, err := service.BatchUpdateLimits(context.Background(), []int64{1}, nil, nil)
require.Error(t, err)
require.Zero(t, affected)
require.Zero(t, repo.calls)
}
func pointerToInt(value int) *int {
return &value
}
@@ -0,0 +1,597 @@
//go:build unit
package service
import (
"context"
"errors"
"net/http"
"reflect"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type accountRepoStubForBulkUpdate struct {
accountRepoStub
bulkUpdateErr error
bulkUpdateIDs []int64
bulkUpdateCalls int
lastBulkUpdate AccountBulkUpdate
bindGroupErrByID map[int64]error
bindGroupsCalls []int64
bindGroupsByAccount map[int64][]int64
createAccount *Account
createID int64
createErr error
updatedAccounts []*Account
updateErr error
getByIDsAccounts []*Account
getByIDsErr error
getByIDsCalled bool
getByIDsIDs []int64
getByIDAccounts map[int64]*Account
getByIDErrByID map[int64]error
getByIDCalled []int64
listByGroupData map[int64][]Account
listByGroupErr map[int64]error
listData []Account
listResult *pagination.PaginationResult
listErr error
listCalled bool
lastListParams pagination.PaginationParams
lastListFilters struct {
platform string
accountType string
status string
search string
groupID int64
privacyMode string
}
}
func (s *accountRepoStubForBulkUpdate) BulkUpdate(_ context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
s.bulkUpdateCalls++
s.bulkUpdateIDs = append([]int64{}, ids...)
s.lastBulkUpdate = updates
if s.bulkUpdateErr != nil {
return 0, s.bulkUpdateErr
}
return int64(len(ids)), nil
}
func requireApplicationErrorReason(t *testing.T, err error, reason string) {
t.Helper()
var appErr *infraerrors.ApplicationError
require.ErrorAs(t, err, &appErr)
require.Equal(t, reason, appErr.Reason)
}
func (s *accountRepoStubForBulkUpdate) Create(_ context.Context, account *Account) error {
s.createAccount = account
if s.createID > 0 {
account.ID = s.createID
}
return s.createErr
}
func (s *accountRepoStubForBulkUpdate) Update(_ context.Context, account *Account) error {
s.updatedAccounts = append(s.updatedAccounts, account)
return s.updateErr
}
func (s *accountRepoStubForBulkUpdate) BindGroups(_ context.Context, accountID int64, groupIDs []int64) error {
s.bindGroupsCalls = append(s.bindGroupsCalls, accountID)
if s.bindGroupsByAccount == nil {
s.bindGroupsByAccount = make(map[int64][]int64)
}
s.bindGroupsByAccount[accountID] = append([]int64{}, groupIDs...)
if err, ok := s.bindGroupErrByID[accountID]; ok {
return err
}
return nil
}
func (s *accountRepoStubForBulkUpdate) GetByIDs(_ context.Context, ids []int64) ([]*Account, error) {
s.getByIDsCalled = true
s.getByIDsIDs = append([]int64{}, ids...)
if s.getByIDsErr != nil {
return nil, s.getByIDsErr
}
return s.getByIDsAccounts, nil
}
func (s *accountRepoStubForBulkUpdate) GetByID(_ context.Context, id int64) (*Account, error) {
s.getByIDCalled = append(s.getByIDCalled, id)
if err, ok := s.getByIDErrByID[id]; ok {
return nil, err
}
if account, ok := s.getByIDAccounts[id]; ok {
return account, nil
}
return nil, errors.New("account not found")
}
func (s *accountRepoStubForBulkUpdate) ListByGroup(_ context.Context, groupID int64) ([]Account, error) {
if err, ok := s.listByGroupErr[groupID]; ok {
return nil, err
}
if rows, ok := s.listByGroupData[groupID]; ok {
return rows, nil
}
return nil, nil
}
func (s *accountRepoStubForBulkUpdate) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) {
return nil, nil
}
func (s *accountRepoStubForBulkUpdate) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
s.listCalled = true
s.lastListParams = params
s.lastListFilters.platform = platform
s.lastListFilters.accountType = accountType
s.lastListFilters.status = status
s.lastListFilters.search = search
s.lastListFilters.groupID = groupID
s.lastListFilters.privacyMode = privacyMode
if s.listErr != nil {
return nil, nil, s.listErr
}
if s.listResult != nil {
return s.listData, s.listResult, nil
}
return s.listData, &pagination.PaginationResult{Total: int64(len(s.listData))}, nil
}
// TestAdminService_BulkUpdateAccounts_AllSuccessIDs 验证批量更新成功时返回 success_ids/failed_ids。
func TestAdminService_BulkUpdateAccounts_AllSuccessIDs(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{}
svc := &adminServiceImpl{accountRepo: repo}
schedulable := true
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2, 3},
Schedulable: &schedulable,
}
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Equal(t, 3, result.Success)
require.Equal(t, 0, result.Failed)
require.ElementsMatch(t, []int64{1, 2, 3}, result.SuccessIDs)
require.Empty(t, result.FailedIDs)
require.Len(t, result.Results, 3)
}
func TestAdminService_BulkUpdateAccounts_RejectsRateChangeForSyncedAccounts(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
getByIDsAccounts: []*Account{
{
ID: 1,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
},
{ID: 2, Extra: map[string]any{}},
},
}
svc := &adminServiceImpl{accountRepo: repo}
rateMultiplier := 0.5
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
RateMultiplier: &rateMultiplier,
})
require.Nil(t, result)
require.Error(t, err)
var appErr *infraerrors.ApplicationError
require.ErrorAs(t, err, &appErr)
require.Equal(t, int32(http.StatusConflict), appErr.Code)
require.Equal(t, "UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT", appErr.Reason)
require.Equal(t, "1", appErr.Metadata["count"])
require.True(t, repo.getByIDsCalled)
require.Empty(t, repo.bulkUpdateIDs, "rate conflict must be rejected before any write")
}
// TestAdminService_BulkUpdateAccounts_PartialFailureIDs 验证部分失败时 success_ids/failed_ids 正确。
func TestAdminService_BulkUpdateAccounts_PartialFailureIDs(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
bindGroupErrByID: map[int64]error{
2: errors.New("bind failed"),
},
}
svc := &adminServiceImpl{
accountRepo: repo,
groupRepo: &groupRepoStubForAdmin{getByID: &Group{ID: 10, Name: "g10"}},
}
groupIDs := []int64{10}
schedulable := false
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2, 3},
GroupIDs: &groupIDs,
Schedulable: &schedulable,
SkipMixedChannelCheck: true,
}
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.Equal(t, 2, result.Success)
require.Equal(t, 1, result.Failed)
require.ElementsMatch(t, []int64{1, 3}, result.SuccessIDs)
require.ElementsMatch(t, []int64{2}, result.FailedIDs)
require.Len(t, result.Results, 3)
}
func TestAdminService_BulkUpdateAccounts_NilGroupRepoReturnsError(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{}
svc := &adminServiceImpl{accountRepo: repo}
groupIDs := []int64{10}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
GroupIDs: &groupIDs,
}
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.Nil(t, result)
require.Error(t, err)
require.Contains(t, err.Error(), "group repository not configured")
}
// TestAdminService_BulkUpdateAccounts_MixedChannelPreCheckBlocksOnExistingConflict verifies
// that the global pre-check detects a conflict with existing group members and returns an
// error before any DB write is performed.
func TestAdminService_BulkUpdateAccounts_MixedChannelPreCheckBlocksOnExistingConflict(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
getByIDsAccounts: []*Account{
{ID: 1, Platform: PlatformAntigravity},
},
// Group 10 already contains an Anthropic account.
listByGroupData: map[int64][]Account{
10: {{ID: 99, Platform: PlatformAnthropic}},
},
}
svc := &adminServiceImpl{
accountRepo: repo,
groupRepo: &groupRepoStubForAdmin{getByID: &Group{ID: 10, Name: "target-group"}},
}
groupIDs := []int64{10}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
GroupIDs: &groupIDs,
}
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.Nil(t, result)
require.Error(t, err)
require.Contains(t, err.Error(), "mixed channel")
// No BindGroups should have been called since the check runs before any write.
require.Empty(t, repo.bindGroupsCalls)
}
func TestAdminServiceBulkUpdateAccounts_ResolvesIDsFromFilters(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
listData: []Account{
{ID: 7},
{ID: 11},
},
listResult: &pagination.PaginationResult{Total: 2},
}
svc := &adminServiceImpl{accountRepo: repo}
schedulable := true
input := &BulkUpdateAccountsInput{
Schedulable: &schedulable,
}
filtersField := reflect.ValueOf(input).Elem().FieldByName("Filters")
require.True(t, filtersField.IsValid(), "BulkUpdateAccountsInput should expose Filters for filter-target bulk update")
require.Equal(t, reflect.Ptr, filtersField.Kind(), "BulkUpdateAccountsInput.Filters should be a pointer field")
filtersValue := reflect.New(filtersField.Type().Elem())
filtersValue.Elem().FieldByName("Platform").SetString(PlatformOpenAI)
filtersValue.Elem().FieldByName("Type").SetString(AccountTypeOAuth)
filtersValue.Elem().FieldByName("Status").SetString(StatusActive)
filtersValue.Elem().FieldByName("Group").SetString("12")
filtersValue.Elem().FieldByName("PrivacyMode").SetString(PrivacyModeCFBlocked)
filtersValue.Elem().FieldByName("Search").SetString("bulk-target")
filtersField.Set(filtersValue)
result, err := svc.BulkUpdateAccounts(context.Background(), input)
require.NoError(t, err)
require.True(t, repo.listCalled, "expected filter-target bulk update to resolve matching IDs via account list filters")
require.Equal(t, PlatformOpenAI, repo.lastListFilters.platform)
require.Equal(t, AccountTypeOAuth, repo.lastListFilters.accountType)
require.Equal(t, StatusActive, repo.lastListFilters.status)
require.Equal(t, "bulk-target", repo.lastListFilters.search)
require.Equal(t, int64(12), repo.lastListFilters.groupID)
require.Equal(t, PrivacyModeCFBlocked, repo.lastListFilters.privacyMode)
require.Equal(t, []int64{7, 11}, repo.bulkUpdateIDs)
require.Equal(t, 2, result.Success)
require.Equal(t, 0, result.Failed)
require.Equal(t, []int64{7, 11}, result.SuccessIDs)
}
func TestAdminServiceBulkUpdateAccounts_NormalizesOpenAISettings(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{
{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
{ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
Credentials: map[string]any{
openAIEndpointCapabilitiesCredentialKey: []any{"chat_completions", "embeddings"},
},
Extra: map[string]any{
openAILongContextBillingEnabledKey: true,
"openai_responses_mode": "auto",
},
})
require.NoError(t, err)
require.Equal(t, 2, result.Success)
require.Zero(t, result.LongContextInheritedCount)
require.Equal(t, 1, repo.bulkUpdateCalls)
require.Contains(t, repo.lastBulkUpdate.Credentials, openAIEndpointCapabilitiesCredentialKey)
require.Nil(t, repo.lastBulkUpdate.Credentials[openAIEndpointCapabilitiesCredentialKey])
require.Equal(t, true, repo.lastBulkUpdate.Extra[openAILongContextBillingEnabledKey])
require.Contains(t, repo.lastBulkUpdate.Extra, "openai_responses_mode")
require.Nil(t, repo.lastBulkUpdate.Extra["openai_responses_mode"])
}
func TestAdminServiceBulkUpdateAccounts_AcceptsLongContextAccountTypes(t *testing.T) {
for _, accountType := range []string{AccountTypeOAuth, AccountTypeSetupToken, AccountTypeAPIKey} {
t.Run(accountType, func(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{{
ID: 1, Platform: PlatformOpenAI, Type: accountType,
}}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{openAILongContextBillingEnabledKey: false},
})
require.NoError(t, err)
require.Equal(t, 1, result.Success)
require.Equal(t, 1, repo.bulkUpdateCalls)
})
}
}
func TestAdminServiceBulkUpdateAccounts_EmbeddingsOnlyResetsResponsesMode(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{
{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
svc := &adminServiceImpl{accountRepo: repo}
_, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: map[string]any{
openAIEndpointCapabilitiesCredentialKey: []string{"embeddings"},
},
})
require.NoError(t, err)
require.Equal(t, []string{"embeddings"}, repo.lastBulkUpdate.Credentials[openAIEndpointCapabilitiesCredentialKey])
require.Contains(t, repo.lastBulkUpdate.Extra, "openai_responses_mode")
require.Nil(t, repo.lastBulkUpdate.Extra["openai_responses_mode"])
}
func TestAdminServiceBulkUpdateAccounts_RejectsInvalidOpenAISettingValuesBeforeWrite(t *testing.T) {
tests := []struct {
name string
credentials map[string]any
extra map[string]any
reason string
}{
{name: "long context type", extra: map[string]any{openAILongContextBillingEnabledKey: "true"}, reason: "OPENAI_LONG_CONTEXT_BILLING_INVALID"},
{name: "empty capabilities", credentials: map[string]any{openAIEndpointCapabilitiesCredentialKey: []any{}}, reason: "OPENAI_ENDPOINT_CAPABILITIES_INVALID"},
{name: "unknown capability", credentials: map[string]any{openAIEndpointCapabilitiesCredentialKey: []any{"responses"}}, reason: "OPENAI_ENDPOINT_CAPABILITIES_INVALID"},
{name: "capabilities type", credentials: map[string]any{openAIEndpointCapabilitiesCredentialKey: "chat_completions"}, reason: "OPENAI_ENDPOINT_CAPABILITIES_INVALID"},
{name: "responses mode", extra: map[string]any{"openai_responses_mode": "sometimes"}, reason: "OPENAI_RESPONSES_MODE_INVALID"},
{name: "responses type", extra: map[string]any{"openai_responses_mode": true}, reason: "OPENAI_RESPONSES_MODE_INVALID"},
{
name: "embeddings conflict",
credentials: map[string]any{openAIEndpointCapabilitiesCredentialKey: []any{"embeddings"}},
extra: map[string]any{"openai_responses_mode": "force_responses"},
reason: "OPENAI_RESPONSES_MODE_INVALID",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: tt.credentials,
Extra: tt.extra,
})
require.Nil(t, result)
requireApplicationErrorReason(t, err, tt.reason)
require.Zero(t, repo.bulkUpdateCalls)
})
}
}
func TestAdminServiceBulkUpdateAccounts_RejectsInvalidOpenAITargetsBeforeWrite(t *testing.T) {
tests := []struct {
name string
accounts []*Account
input *BulkUpdateAccountsInput
}{
{
name: "missing account",
accounts: []*Account{{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth}},
input: &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
},
},
{
name: "mixed platform long context",
accounts: []*Account{{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth}},
input: &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
},
},
{
name: "oauth endpoint capabilities",
accounts: []*Account{{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth}},
input: &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: map[string]any{openAIEndpointCapabilitiesCredentialKey: nil},
},
},
{
name: "unsupported OpenAI long context account type",
accounts: []*Account{{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeServiceAccount}},
input: &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: tt.accounts}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), tt.input)
require.Nil(t, result)
requireApplicationErrorReason(t, err, "OPENAI_BULK_TARGET_INVALID")
require.Zero(t, repo.bulkUpdateCalls)
})
}
}
func TestAdminServiceBulkUpdateAccounts_ForcedResponsesRequiresChatCapability(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
openAIEndpointCapabilitiesCredentialKey: []any{"embeddings"},
},
}}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{"openai_responses_mode": "force_chat_completions"},
})
require.Nil(t, result)
requireApplicationErrorReason(t, err, "OPENAI_BULK_TARGET_INVALID")
require.Zero(t, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccounts_ForcedResponsesAcceptsChatCapabilityUpdate(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
openAIEndpointCapabilitiesCredentialKey: []any{"embeddings"},
},
}}}
svc := &adminServiceImpl{accountRepo: repo}
_, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Credentials: map[string]any{
openAIEndpointCapabilitiesCredentialKey: []any{"chat_completions"},
},
Extra: map[string]any{"openai_responses_mode": "force_responses"},
})
require.NoError(t, err)
require.Equal(t, 1, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccounts_ReportsLongContextShadowInheritance(t *testing.T) {
parentID := int64(1)
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{
{ID: parentID, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
{ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID},
}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{parentID, 2},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
})
require.NoError(t, err)
require.Equal(t, 1, result.LongContextInheritedCount)
require.Equal(t, 1, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccounts_RequiresParentForShadowOnlyLongContextUpdate(t *testing.T) {
parentID := int64(10)
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{
{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID},
{ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID},
}}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
})
require.Nil(t, result)
requireApplicationErrorReason(t, err, "OPENAI_LONG_CONTEXT_PARENT_REQUIRED")
require.Zero(t, repo.bulkUpdateCalls)
}
func TestAdminServiceBulkUpdateAccounts_ShadowLongContextAllowsOtherUpdates(t *testing.T) {
parentID := int64(10)
repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{{
ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID,
}}}
svc := &adminServiceImpl{accountRepo: repo}
status := StatusDisabled
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Status: status,
Extra: map[string]any{openAILongContextBillingEnabledKey: false},
})
require.NoError(t, err)
require.Equal(t, 1, result.LongContextInheritedCount)
require.Equal(t, 1, repo.bulkUpdateCalls)
require.NotNil(t, repo.lastBulkUpdate.Status)
require.Equal(t, status, *repo.lastBulkUpdate.Status)
}
func TestAdminServiceBulkUpdateAccounts_ValidatesFilterResolvedOpenAITargets(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
listData: []Account{{ID: 7}},
listResult: &pagination.PaginationResult{Total: 1},
getByIDsAccounts: []*Account{{ID: 7, Platform: PlatformAnthropic, Type: AccountTypeOAuth}},
}
svc := &adminServiceImpl{accountRepo: repo}
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
Filters: &BulkUpdateAccountFilters{Platform: PlatformOpenAI},
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
})
require.Nil(t, result)
requireApplicationErrorReason(t, err, "OPENAI_BULK_TARGET_INVALID")
require.Equal(t, []int64{7}, repo.getByIDsIDs)
require.Zero(t, repo.bulkUpdateCalls)
}
@@ -0,0 +1,88 @@
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type accountRepoStubForClearAccountError struct {
mockAccountRepoForGemini
account *Account
clearErrorCalls int
clearRateLimitCalls int
clearAntigravityCalls int
clearModelRateLimitCalls int
clearTempUnschedCalls int
}
func (r *accountRepoStubForClearAccountError) GetByID(ctx context.Context, id int64) (*Account, error) {
return r.account, nil
}
func (r *accountRepoStubForClearAccountError) ClearError(ctx context.Context, id int64) error {
r.clearErrorCalls++
r.account.Status = StatusActive
r.account.ErrorMessage = ""
return nil
}
func (r *accountRepoStubForClearAccountError) ClearRateLimit(ctx context.Context, id int64) error {
r.clearRateLimitCalls++
r.account.RateLimitedAt = nil
r.account.RateLimitResetAt = nil
return nil
}
func (r *accountRepoStubForClearAccountError) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
r.clearAntigravityCalls++
return nil
}
func (r *accountRepoStubForClearAccountError) ClearModelRateLimits(ctx context.Context, id int64) error {
r.clearModelRateLimitCalls++
return nil
}
func (r *accountRepoStubForClearAccountError) ClearTempUnschedulable(ctx context.Context, id int64) error {
r.clearTempUnschedCalls++
r.account.TempUnschedulableUntil = nil
r.account.TempUnschedulableReason = ""
return nil
}
func TestAdminService_ClearAccountError_AlsoClearsRecoverableRuntimeState(t *testing.T) {
until := time.Now().Add(10 * time.Minute)
resetAt := time.Now().Add(5 * time.Minute)
repo := &accountRepoStubForClearAccountError{
account: &Account{
ID: 31,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusError,
ErrorMessage: "refresh failed",
RateLimitResetAt: &resetAt,
TempUnschedulableUntil: &until,
TempUnschedulableReason: "missing refresh token",
},
}
blocker := &runtimeBlockRecorder{}
svc := &adminServiceImpl{accountRepo: repo, runtimeBlocker: blocker}
updated, err := svc.ClearAccountError(context.Background(), 31)
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.clearErrorCalls)
require.Equal(t, 1, repo.clearRateLimitCalls)
require.Equal(t, 1, repo.clearAntigravityCalls)
require.Equal(t, 1, repo.clearModelRateLimitCalls)
require.Equal(t, 1, repo.clearTempUnschedCalls)
require.Nil(t, updated.RateLimitResetAt)
require.Nil(t, updated.TempUnschedulableUntil)
require.Empty(t, updated.TempUnschedulableReason)
require.Equal(t, []int64{31}, blocker.clearedIDs)
}
@@ -0,0 +1,214 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/stretchr/testify/require"
)
type accountRepoStubForCompositeModelsList struct {
accountRepoStub
accounts []Account
}
func (s *accountRepoStubForCompositeModelsList) ListSchedulableByGroupID(_ context.Context, _ int64) ([]Account, error) {
return s.accounts, nil
}
func TestAdminService_CreateCompositeGroupCopiesAccountsFromConcreteGroups(t *testing.T) {
var copiedFrom []int64
var boundGroupID int64
var boundAccountIDs []int64
groupRepo := &groupRepoStubForAdmin{
createID: 99,
getByIDByID: map[int64]*Group{
10: {ID: 10, Platform: PlatformOpenAI},
20: {ID: 20, Platform: PlatformGemini},
},
getAccountIDsByGroupIDsFn: func(groupIDs []int64) ([]int64, error) {
copiedFrom = append([]int64{}, groupIDs...)
return []int64{101, 202}, nil
},
bindAccountsToGroupFn: func(groupID int64, accountIDs []int64) error {
boundGroupID = groupID
boundAccountIDs = append([]int64{}, accountIDs...)
return nil
},
}
svc := &adminServiceImpl{groupRepo: groupRepo}
group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "Composite",
Platform: PlatformComposite,
RateMultiplier: 1,
MaxReasoningEffort: "medium",
ReasoningEffortMappings: []ReasoningEffortMapping{
{From: "max", To: "xhigh"},
},
CopyAccountsFromGroupIDs: []int64{10, 20, 10},
})
require.NoError(t, err)
require.Equal(t, PlatformComposite, groupRepo.created.Platform)
require.Equal(t, "medium", groupRepo.created.MaxReasoningEffort)
require.Equal(t, []ReasoningEffortMapping{{From: "max", To: "xhigh"}}, groupRepo.created.ReasoningEffortMappings)
require.Equal(t, int64(99), group.ID)
require.Equal(t, int64(2), group.AccountCount)
require.ElementsMatch(t, []int64{10, 20}, copiedFrom)
require.Equal(t, int64(99), boundGroupID)
require.ElementsMatch(t, []int64{101, 202}, boundAccountIDs)
}
func TestAdminService_UpdateCompositeGroupCopiesAccountsFromConcreteGroups(t *testing.T) {
var clearedGroupID int64
var copiedFrom []int64
var boundGroupID int64
var boundAccountIDs []int64
groupRepo := &groupRepoStubForAdmin{
getByIDByID: map[int64]*Group{
10: {ID: 10, Platform: PlatformOpenAI},
20: {ID: 20, Platform: PlatformGrok},
99: {ID: 99, Platform: PlatformComposite, RateMultiplier: 1, SubscriptionType: SubscriptionTypeStandard},
},
deleteAccountGroupsByGroupIDFn: func(groupID int64) (int64, error) {
clearedGroupID = groupID
return 2, nil
},
getAccountIDsByGroupIDsFn: func(groupIDs []int64) ([]int64, error) {
copiedFrom = append([]int64{}, groupIDs...)
return []int64{301, 302}, nil
},
bindAccountsToGroupFn: func(groupID int64, accountIDs []int64) error {
boundGroupID = groupID
boundAccountIDs = append([]int64{}, accountIDs...)
return nil
},
}
svc := &adminServiceImpl{groupRepo: groupRepo}
maxReasoningEffort := "low"
reasoningEffortMappings := []ReasoningEffortMapping{{From: "max", To: "high"}}
group, err := svc.UpdateGroup(context.Background(), 99, &UpdateGroupInput{
MaxReasoningEffort: &maxReasoningEffort,
ReasoningEffortMappings: &reasoningEffortMappings,
CopyAccountsFromGroupIDs: []int64{10, 20},
})
require.NoError(t, err)
require.Equal(t, PlatformComposite, group.Platform)
require.Equal(t, "low", group.MaxReasoningEffort)
require.Equal(t, reasoningEffortMappings, group.ReasoningEffortMappings)
require.Equal(t, int64(99), clearedGroupID)
require.ElementsMatch(t, []int64{10, 20}, copiedFrom)
require.Equal(t, int64(99), boundGroupID)
require.ElementsMatch(t, []int64{301, 302}, boundAccountIDs)
}
func TestAdminService_CreateAccountAllowsCompositeGroupAssignment(t *testing.T) {
accountRepo := &accountRepoStubForBulkUpdate{createID: 7}
groupRepo := &groupRepoStubForAdmin{
getByIDByID: map[int64]*Group{
99: {ID: 99, Platform: PlatformComposite},
},
}
svc := &adminServiceImpl{accountRepo: accountRepo, groupRepo: groupRepo}
account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
Name: "OpenAI account",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
GroupIDs: []int64{99},
SkipDefaultGroupBind: true,
SkipMixedChannelCheck: true,
})
require.NoError(t, err)
require.Equal(t, int64(7), account.ID)
require.Equal(t, PlatformOpenAI, accountRepo.createAccount.Platform)
require.ElementsMatch(t, []int64{99}, accountRepo.bindGroupsByAccount[7])
}
func TestAdminService_UpdateAccountAllowsCompositeGroupAssignment(t *testing.T) {
accountRepo := &accountRepoStubForBulkUpdate{
getByIDAccounts: map[int64]*Account{
7: {ID: 7, Platform: PlatformGemini, Type: AccountTypeAPIKey, Status: StatusActive, Extra: map[string]any{}},
},
}
groupRepo := &groupRepoStubForAdmin{
getByIDByID: map[int64]*Group{
99: {ID: 99, Platform: PlatformComposite},
},
}
svc := &adminServiceImpl{accountRepo: accountRepo, groupRepo: groupRepo}
groupIDs := []int64{99}
account, err := svc.UpdateAccount(context.Background(), 7, &UpdateAccountInput{
GroupIDs: &groupIDs,
SkipMixedChannelCheck: true,
})
require.NoError(t, err)
require.Equal(t, int64(7), account.ID)
require.Len(t, accountRepo.updatedAccounts, 1)
require.ElementsMatch(t, []int64{99}, accountRepo.bindGroupsByAccount[7])
}
func TestAdminService_CompositeModelsListCandidatesIncludeConcreteAccountMappings(t *testing.T) {
accountRepo := &accountRepoStubForCompositeModelsList{
accounts: []Account{
{
ID: 1,
Platform: PlatformOpenAI,
Credentials: map[string]any{
"model_mapping": map[string]any{"gpt-custom": "gpt-5"},
},
},
{
ID: 2,
Platform: PlatformGemini,
Credentials: map[string]any{
"model_mapping": map[string]any{"gemini-custom": "gemini-2.5-flash"},
},
},
{
ID: 3,
Platform: PlatformKimi,
Credentials: map[string]any{
"model_mapping": map[string]any{"kimi-custom": "kimi-k2"},
},
},
},
}
groupRepo := &groupRepoStubForAdmin{
getByIDByID: map[int64]*Group{
99: {ID: 99, Platform: PlatformComposite},
},
}
svc := &adminServiceImpl{accountRepo: accountRepo, groupRepo: groupRepo}
candidates, err := svc.GetGroupModelsListCandidates(context.Background(), 99, PlatformComposite)
require.NoError(t, err)
require.Contains(t, candidates, "gpt-custom")
require.Contains(t, candidates, "gemini-custom")
require.Contains(t, candidates, "kimi-custom")
require.Contains(t, candidates, "gpt-5.5")
require.Contains(t, candidates, "gemini-2.5-flash")
}
// 独立 CN 分组的模型列表候选沿用 default 分支的 Claude 默认列表;
// composite 支持不得改变独立分组的候选语义。
func TestAdminService_CNProviderModelsListCandidatesKeepClaudeDefaults(t *testing.T) {
want := make([]string, 0, len(claude.DefaultModels))
for _, model := range claude.DefaultModels {
want = append(want, model.ID)
}
for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek} {
require.Equal(t, want, defaultModelsListCandidateIDs(platform), "platform=%s", platform)
}
}
@@ -0,0 +1,148 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
func TestAdminService_CreateUser_Success(t *testing.T) {
repo := &userRepoStub{nextID: 10}
svc := &adminServiceImpl{userRepo: repo}
balance := 12.5
input := &CreateUserInput{
Email: "user@test.com",
Password: "strong-pass",
Username: "tester",
Notes: "note",
Balance: &balance,
Concurrency: 7,
AllowedGroups: []int64{3, 5},
}
user, err := svc.CreateUser(context.Background(), input)
require.NoError(t, err)
require.NotNil(t, user)
require.Equal(t, int64(10), user.ID)
require.Equal(t, input.Email, user.Email)
require.Equal(t, input.Username, user.Username)
require.Equal(t, input.Notes, user.Notes)
require.Equal(t, balance, user.Balance)
require.Equal(t, input.Concurrency, user.Concurrency)
require.Equal(t, input.AllowedGroups, user.AllowedGroups)
require.Equal(t, RoleUser, user.Role)
require.Equal(t, StatusActive, user.Status)
require.True(t, user.CheckPassword(input.Password))
require.Len(t, repo.created, 1)
require.Equal(t, user, repo.created[0])
}
func TestAdminService_CreateUser_UsesDefaultBalanceWhenBalanceOmitted(t *testing.T) {
repo := &userRepoStub{nextID: 11}
cfg := &config.Config{
Default: config.DefaultConfig{
UserBalance: 0,
},
}
settingService := NewSettingService(&settingRepoStub{values: map[string]string{
SettingKeyDefaultBalance: "0.02",
}}, cfg)
svc := &adminServiceImpl{userRepo: repo, settingService: settingService}
user, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "default-balance@test.com",
Password: "strong-pass",
})
require.NoError(t, err)
require.NotNil(t, user)
require.Equal(t, 0.02, user.Balance)
require.Len(t, repo.created, 1)
require.Equal(t, 0.02, repo.created[0].Balance)
}
func TestAdminService_CreateUser_ExplicitZeroBalanceOverridesDefault(t *testing.T) {
repo := &userRepoStub{nextID: 12}
cfg := &config.Config{
Default: config.DefaultConfig{
UserBalance: 0,
},
}
settingService := NewSettingService(&settingRepoStub{values: map[string]string{
SettingKeyDefaultBalance: "0.02",
}}, cfg)
svc := &adminServiceImpl{userRepo: repo, settingService: settingService}
balance := 0.0
user, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "zero-balance@test.com",
Password: "strong-pass",
Balance: &balance,
})
require.NoError(t, err)
require.NotNil(t, user)
require.Equal(t, 0.0, user.Balance)
require.Len(t, repo.created, 1)
require.Equal(t, 0.0, repo.created[0].Balance)
}
func TestAdminService_CreateUser_EmailExists(t *testing.T) {
repo := &userRepoStub{createErr: ErrEmailExists}
svc := &adminServiceImpl{userRepo: repo}
_, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "dup@test.com",
Password: "password",
})
require.ErrorIs(t, err, ErrEmailExists)
require.Empty(t, repo.created)
}
func TestAdminService_CreateUser_CreateError(t *testing.T) {
createErr := errors.New("db down")
repo := &userRepoStub{createErr: createErr}
svc := &adminServiceImpl{userRepo: repo}
_, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "user@test.com",
Password: "password",
})
require.ErrorIs(t, err, createErr)
require.Empty(t, repo.created)
}
func TestAdminService_CreateUser_AssignsDefaultSubscriptions(t *testing.T) {
repo := &userRepoStub{nextID: 21}
assigner := &defaultSubscriptionAssignerStub{}
cfg := &config.Config{
Default: config.DefaultConfig{
UserBalance: 0,
UserConcurrency: 1,
},
}
settingService := NewSettingService(&settingRepoStub{values: map[string]string{
SettingKeyDefaultSubscriptions: `[{"group_id":5,"validity_days":30}]`,
}}, cfg)
svc := &adminServiceImpl{
userRepo: repo,
settingService: settingService,
defaultSubAssigner: assigner,
}
_, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "new-user@test.com",
Password: "password",
})
require.NoError(t, err)
require.Len(t, assigner.calls, 1)
require.Equal(t, int64(21), assigner.calls[0].UserID)
require.Equal(t, int64(5), assigner.calls[0].GroupID)
require.Equal(t, 30, assigner.calls[0].ValidityDays)
}
@@ -0,0 +1,117 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
type updateAccountCredsRepoStub struct {
mockAccountRepoForGemini
account *Account
updateCalls int
}
func (r *updateAccountCredsRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
return r.account, nil
}
func (r *updateAccountCredsRepoStub) Update(ctx context.Context, account *Account) error {
r.updateCalls++
r.account = account
return nil
}
func TestUpdateAccount_PreservesSensitiveCredsWhenIncomingOmits(t *testing.T) {
accountID := int64(202)
repo := &updateAccountCredsRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"refresh_token": "rt-existing",
"access_token": "at-existing",
"id_token": "id-existing",
"base_url": "https://old.example.com",
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
// 模拟前端编辑:仅修改 base_url,没有传 token(脱敏后前端 spread 拿不到敏感键)
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Credentials: map[string]any{
"base_url": "https://new.example.com",
},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
// 敏感键应保留
require.Equal(t, "rt-existing", repo.account.Credentials["refresh_token"])
require.Equal(t, "at-existing", repo.account.Credentials["access_token"])
require.Equal(t, "id-existing", repo.account.Credentials["id_token"])
// 非敏感键被替换
require.Equal(t, "https://new.example.com", repo.account.Credentials["base_url"])
}
func TestUpdateAccount_ExplicitNewTokenOverwrites(t *testing.T) {
accountID := int64(203)
repo := &updateAccountCredsRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"refresh_token": "rt-old",
"api_key": "sk-old",
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Credentials: map[string]any{
"refresh_token": "rt-new",
// api_key 没传 → 应保留旧值
},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, "rt-new", repo.account.Credentials["refresh_token"])
require.Equal(t, "sk-old", repo.account.Credentials["api_key"])
}
func TestUpdateAccount_EmptyCredentialsSkipsUpdate(t *testing.T) {
accountID := int64(204)
repo := &updateAccountCredsRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"refresh_token": "rt-existing",
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
_, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Credentials: map[string]any{}, // len == 0 → 闸门跳过
Name: "renamed",
})
require.NoError(t, err)
require.Equal(t, "rt-existing", repo.account.Credentials["refresh_token"], "空 credentials 不应触碰已有 token")
require.Equal(t, "renamed", repo.account.Name)
}
@@ -0,0 +1,792 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type userRepoStub struct {
user *User
usersByID map[int64]*User
getErr error
createErr error
deleteErr error
exists bool
existsErr error
aliasExists bool
aliasErr error
guardedCreates int
nextID int64
created []*User
updated []*User
deletedIDs []int64
usersByEmail map[string]*User
getByEmailErr error
getByEmailMisses int
domainCounts map[string]int
domainCountErr error
domainLimitErr error
domainLimitedCreates int
}
func (s *userRepoStub) CountUsersByEmailDomain(_ context.Context, domain string) (int, error) {
if s.domainCountErr != nil {
return 0, s.domainCountErr
}
return s.domainCounts[domain], nil
}
func (s *userRepoStub) CreateWithEmailAliasGuardAndDomainLimit(ctx context.Context, user *User, domain string) error {
s.domainLimitedCreates++
if s.domainLimitErr != nil {
return s.domainLimitErr
}
if s.domainCounts[domain] > 0 {
return ErrEmailDomainRegistrationLimit
}
return s.CreateWithEmailAliasGuard(ctx, user)
}
func (s *userRepoStub) Create(ctx context.Context, user *User) error {
if s.createErr != nil {
return s.createErr
}
if s.nextID != 0 && user.ID == 0 {
user.ID = s.nextID
}
s.created = append(s.created, user)
if s.usersByEmail == nil {
s.usersByEmail = make(map[string]*User)
}
s.usersByEmail[user.Email] = user
s.user = user
return nil
}
func (s *userRepoStub) CreateWithEmailAliasGuard(ctx context.Context, user *User) error {
s.guardedCreates++
if s.aliasErr != nil {
return s.aliasErr
}
if s.aliasExists {
return ErrEmailExists
}
return s.Create(ctx, user)
}
func (s *userRepoStub) GetByID(ctx context.Context, id int64) (*User, error) {
if s.getErr != nil {
return nil, s.getErr
}
if s.usersByID != nil {
if user, ok := s.usersByID[id]; ok {
return user, nil
}
return nil, ErrUserNotFound
}
if s.user == nil {
return nil, ErrUserNotFound
}
return s.user, nil
}
func (s *userRepoStub) GetByEmail(ctx context.Context, email string) (*User, error) {
if s.getByEmailErr != nil {
return nil, s.getByEmailErr
}
if s.getByEmailMisses > 0 {
s.getByEmailMisses--
return nil, ErrUserNotFound
}
if s.usersByEmail != nil {
if user, ok := s.usersByEmail[email]; ok {
return user, nil
}
}
if s.user != nil && s.user.Email == email {
return s.user, nil
}
return nil, ErrUserNotFound
}
func (s *userRepoStub) GetFirstAdmin(ctx context.Context) (*User, error) {
panic("unexpected GetFirstAdmin call")
}
func (s *userRepoStub) Update(ctx context.Context, user *User, fields UserUpdateFields) error {
s.updated = append(s.updated, user)
if s.usersByEmail == nil {
s.usersByEmail = make(map[string]*User)
}
s.usersByEmail[user.Email] = user
s.user = user
return nil
}
func (s *userRepoStub) Delete(ctx context.Context, id int64) error {
s.deletedIDs = append(s.deletedIDs, id)
return s.deleteErr
}
func (s *userRepoStub) GetUserAvatar(ctx context.Context, userID int64) (*UserAvatar, error) {
panic("unexpected GetUserAvatar call")
}
func (s *userRepoStub) UpsertUserAvatar(ctx context.Context, userID int64, input UpsertUserAvatarInput) (*UserAvatar, error) {
panic("unexpected UpsertUserAvatar call")
}
func (s *userRepoStub) DeleteUserAvatar(ctx context.Context, userID int64) error {
panic("unexpected DeleteUserAvatar call")
}
func (s *userRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]User, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *userRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters UserListFilters) ([]User, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *userRepoStub) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserIDs call")
}
func (s *userRepoStub) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserID call")
}
func (s *userRepoStub) UpdateUserLastActiveAt(ctx context.Context, userID int64, activeAt time.Time) error {
panic("unexpected UpdateUserLastActiveAt call")
}
func (s *userRepoStub) UpdateBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected UpdateBalance call")
}
func (s *userRepoStub) DeductBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected DeductBalance call")
}
func (s *userRepoStub) AdjustBalance(ctx context.Context, id int64, delta float64) (BalanceChange, error) {
panic("unexpected AdjustBalance call")
}
func (s *userRepoStub) SetBalance(ctx context.Context, id int64, value float64) (BalanceChange, error) {
panic("unexpected SetBalance call")
}
func (s *userRepoStub) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
panic("unexpected UpdateConcurrency call")
}
func (s *userRepoStub) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *userRepoStub) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *userRepoStub) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
return 0, nil
}
func (s *userRepoStub) ExistsByEmail(ctx context.Context, email string) (bool, error) {
if s.existsErr != nil {
return false, s.existsErr
}
return s.exists, nil
}
func (s *userRepoStub) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
if s.aliasErr != nil {
return false, s.aliasErr
}
return s.aliasExists, nil
}
func (s *userRepoStub) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
func (s *userRepoStub) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected RemoveGroupFromUserAllowedGroups call")
}
func (s *userRepoStub) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected AddGroupToAllowedGroups call")
}
func (s *userRepoStub) ListUserAuthIdentities(ctx context.Context, userID int64) ([]UserAuthIdentityRecord, error) {
panic("unexpected ListUserAuthIdentities call")
}
func (s *userRepoStub) UnbindUserAuthProvider(context.Context, int64, string) error {
panic("unexpected UnbindUserAuthProvider call")
}
func (s *userRepoStub) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
panic("unexpected UpdateTotpSecret call")
}
func (s *userRepoStub) EnableTotp(ctx context.Context, userID int64) error {
panic("unexpected EnableTotp call")
}
func (s *userRepoStub) DisableTotp(ctx context.Context, userID int64) error {
panic("unexpected DisableTotp call")
}
func (s *userRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) {
return s.GetByID(ctx, id)
}
type groupRepoStub struct {
affectedUserIDs []int64
deleteErr error
deleteCalls []int64
}
func (s *groupRepoStub) Create(ctx context.Context, group *Group) error {
panic("unexpected Create call")
}
func (s *groupRepoStub) GetByID(ctx context.Context, id int64) (*Group, error) {
panic("unexpected GetByID call")
}
func (s *groupRepoStub) GetByIDLite(ctx context.Context, id int64) (*Group, error) {
panic("unexpected GetByIDLite call")
}
func (s *groupRepoStub) Update(ctx context.Context, group *Group) error {
panic("unexpected Update call")
}
func (s *groupRepoStub) Delete(ctx context.Context, id int64) error {
panic("unexpected Delete call")
}
func (s *groupRepoStub) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
s.deleteCalls = append(s.deleteCalls, id)
return s.affectedUserIDs, s.deleteErr
}
func (s *groupRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *groupRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *groupRepoStub) ListActive(ctx context.Context) ([]Group, error) {
panic("unexpected ListActive call")
}
func (s *groupRepoStub) ListActiveByPlatform(ctx context.Context, platform string) ([]Group, error) {
panic("unexpected ListActiveByPlatform call")
}
func (s *groupRepoStub) ExistsByName(ctx context.Context, name string) (bool, error) {
panic("unexpected ExistsByName call")
}
func (s *groupRepoStub) GetAccountCount(ctx context.Context, groupID int64) (int64, int64, error) {
panic("unexpected GetAccountCount call")
}
func (s *groupRepoStub) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected DeleteAccountGroupsByGroupID call")
}
func (s *groupRepoStub) BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error {
panic("unexpected BindAccountsToGroup call")
}
func (s *groupRepoStub) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error) {
panic("unexpected GetAccountIDsByGroupIDs call")
}
func (s *groupRepoStub) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
return nil
}
type deleteGroupAPIKeyRepoStub struct {
apiKeyRepoStubForGroupUpdate
keys []string
listErr error
listGroupIDs []int64
}
func (s *deleteGroupAPIKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
s.listGroupIDs = append(s.listGroupIDs, groupID)
if s.listErr != nil {
return nil, s.listErr
}
return s.keys, nil
}
type proxyRepoStub struct {
deleteErr error
countErr error
accountCount int64
deletedIDs []int64
}
func (s *proxyRepoStub) Create(ctx context.Context, proxy *Proxy) error {
panic("unexpected Create call")
}
func (s *proxyRepoStub) GetByID(ctx context.Context, id int64) (*Proxy, error) {
panic("unexpected GetByID call")
}
func (s *proxyRepoStub) ListByIDs(ctx context.Context, ids []int64) ([]Proxy, error) {
panic("unexpected ListByIDs call")
}
func (s *proxyRepoStub) Update(ctx context.Context, proxy *Proxy) error {
panic("unexpected Update call")
}
func (s *proxyRepoStub) Delete(ctx context.Context, id int64) error {
s.deletedIDs = append(s.deletedIDs, id)
return s.deleteErr
}
func (s *proxyRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]Proxy, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *proxyRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]Proxy, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *proxyRepoStub) ListActive(ctx context.Context) ([]Proxy, error) {
panic("unexpected ListActive call")
}
func (s *proxyRepoStub) ListActiveWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error) {
panic("unexpected ListActiveWithAccountCount call")
}
func (s *proxyRepoStub) ListWithFiltersAndAccountCount(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]ProxyWithAccountCount, *pagination.PaginationResult, error) {
panic("unexpected ListWithFiltersAndAccountCount call")
}
func (s *proxyRepoStub) ExistsByHostPortAuth(ctx context.Context, host string, port int, username, password string) (bool, error) {
panic("unexpected ExistsByHostPortAuth call")
}
func (s *proxyRepoStub) CountAccountsByProxyID(ctx context.Context, proxyID int64) (int64, error) {
if s.countErr != nil {
return 0, s.countErr
}
return s.accountCount, nil
}
func (s *proxyRepoStub) ListAccountSummariesByProxyID(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error) {
panic("unexpected ListAccountSummariesByProxyID call")
}
func (s *proxyRepoStub) SweepExpiredProxies(_ context.Context, _ time.Time) (int64, error) {
return 0, nil
}
func (s *proxyRepoStub) ListAllForFallback(_ context.Context) ([]Proxy, error) {
return nil, nil
}
func (s *proxyRepoStub) CountExpired(_ context.Context) (int64, error) {
return 0, nil
}
func (s *proxyRepoStub) CountExpiringSoon(_ context.Context, _ time.Time) (int64, error) {
return 0, nil
}
type redeemRepoStub struct {
deleteErrByID map[int64]error
deletedIDs []int64
batchUpdateIDs []int64
batchUpdateFields RedeemCodeBatchUpdateFields
batchUpdateResult int64
batchUpdateErr error
batchUpdateCalled bool
}
func (s *redeemRepoStub) Create(ctx context.Context, code *RedeemCode) error {
panic("unexpected Create call")
}
func (s *redeemRepoStub) CreateBatch(ctx context.Context, codes []RedeemCode) error {
panic("unexpected CreateBatch call")
}
func (s *redeemRepoStub) GetByID(ctx context.Context, id int64) (*RedeemCode, error) {
panic("unexpected GetByID call")
}
func (s *redeemRepoStub) GetByCode(ctx context.Context, code string) (*RedeemCode, error) {
panic("unexpected GetByCode call")
}
func (s *redeemRepoStub) Update(ctx context.Context, code *RedeemCode) error {
panic("unexpected Update call")
}
func (s *redeemRepoStub) BatchUpdate(ctx context.Context, ids []int64, fields RedeemCodeBatchUpdateFields) (int64, error) {
s.batchUpdateCalled = true
s.batchUpdateIDs = append([]int64(nil), ids...)
s.batchUpdateFields = fields
if s.batchUpdateErr != nil {
return 0, s.batchUpdateErr
}
if s.batchUpdateResult != 0 {
return s.batchUpdateResult, nil
}
return int64(len(ids)), nil
}
func (s *redeemRepoStub) Delete(ctx context.Context, id int64) error {
s.deletedIDs = append(s.deletedIDs, id)
if s.deleteErrByID != nil {
if err, ok := s.deleteErrByID[id]; ok {
return err
}
}
return nil
}
func (s *redeemRepoStub) Use(ctx context.Context, id, userID int64) error {
panic("unexpected Use call")
}
func (s *redeemRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]RedeemCode, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *redeemRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, codeType, status, search string) ([]RedeemCode, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *redeemRepoStub) ListByUser(ctx context.Context, userID int64, limit int) ([]RedeemCode, error) {
panic("unexpected ListByUser call")
}
func (s *redeemRepoStub) ListByUserPaginated(ctx context.Context, userID int64, params pagination.PaginationParams, codeType string) ([]RedeemCode, *pagination.PaginationResult, error) {
panic("unexpected ListByUserPaginated call")
}
func (s *redeemRepoStub) SumPositiveBalanceByUser(ctx context.Context, userID int64) (float64, error) {
panic("unexpected SumPositiveBalanceByUser call")
}
type subscriptionInvalidateCall struct {
userID int64
groupID int64
}
type billingCacheStub struct {
invalidations chan subscriptionInvalidateCall
}
func newBillingCacheStub(buffer int) *billingCacheStub {
return &billingCacheStub{invalidations: make(chan subscriptionInvalidateCall, buffer)}
}
func (s *billingCacheStub) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
panic("unexpected GetUserBalance call")
}
func (s *billingCacheStub) SetUserBalance(ctx context.Context, userID int64, balance float64) error {
panic("unexpected SetUserBalance call")
}
func (s *billingCacheStub) DeductUserBalance(ctx context.Context, userID int64, amount float64) error {
panic("unexpected DeductUserBalance call")
}
func (s *billingCacheStub) InvalidateUserBalance(ctx context.Context, userID int64) error {
panic("unexpected InvalidateUserBalance call")
}
func (s *billingCacheStub) GetSubscriptionCache(ctx context.Context, userID, groupID int64) (*SubscriptionCacheData, error) {
panic("unexpected GetSubscriptionCache call")
}
func (s *billingCacheStub) SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *SubscriptionCacheData) error {
panic("unexpected SetSubscriptionCache call")
}
func (s *billingCacheStub) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error {
panic("unexpected UpdateSubscriptionUsage call")
}
func (s *billingCacheStub) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error {
s.invalidations <- subscriptionInvalidateCall{userID: userID, groupID: groupID}
return nil
}
func (s *billingCacheStub) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*APIKeyRateLimitCacheData, error) {
panic("unexpected GetAPIKeyRateLimit call")
}
func (s *billingCacheStub) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *APIKeyRateLimitCacheData) error {
panic("unexpected SetAPIKeyRateLimit call")
}
func (s *billingCacheStub) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
panic("unexpected UpdateAPIKeyRateLimitUsage call")
}
func (s *billingCacheStub) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
panic("unexpected InvalidateAPIKeyRateLimit call")
}
func (s *billingCacheStub) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaCacheEntry, bool, error) {
panic("unexpected GetUserPlatformQuotaCache call")
}
func (s *billingCacheStub) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error {
panic("unexpected SetUserPlatformQuotaCache call")
}
func (s *billingCacheStub) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error {
panic("unexpected DeleteUserPlatformQuotaCache call")
}
func (s *billingCacheStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error {
panic("unexpected IncrUserPlatformQuotaUsageCache call")
}
func (s *billingCacheStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) {
panic("unexpected PopDirtyUserPlatformQuotaKeys call")
}
func (s *billingCacheStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error {
panic("unexpected ReaddDirtyUserPlatformQuotaKeys call")
}
func (s *billingCacheStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) {
panic("unexpected BatchGetUserPlatformQuotaCache call")
}
func waitForInvalidations(t *testing.T, ch <-chan subscriptionInvalidateCall, expected int) []subscriptionInvalidateCall {
t.Helper()
calls := make([]subscriptionInvalidateCall, 0, expected)
timeout := time.After(2 * time.Second)
for len(calls) < expected {
select {
case call := <-ch:
calls = append(calls, call)
case <-timeout:
t.Fatalf("timeout waiting for %d invalidations, got %d", expected, len(calls))
}
}
return calls
}
func TestAdminService_DeleteUser_Success(t *testing.T) {
repo := &userRepoStub{user: &User{ID: 7, Role: RoleUser}}
svc := &adminServiceImpl{userRepo: repo}
err := svc.DeleteUser(context.Background(), 7)
require.NoError(t, err)
require.Equal(t, []int64{7}, repo.deletedIDs)
}
func TestAdminService_DeleteUser_DeletesOwnedAPIKeys(t *testing.T) {
repo := &userRepoStub{user: &User{ID: 7, Role: RoleUser}}
apiKeyRepo := &apiKeyRepoStub{
allowListByUserID: true,
listByUserIDKeys: []APIKey{
{ID: 11, UserID: 7, Key: "sk-user-1"},
{ID: 12, UserID: 7, Key: "sk-user-2"},
},
}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
apiKeyRepo: apiKeyRepo,
authCacheInvalidator: invalidator,
}
err := svc.DeleteUser(context.Background(), 7)
require.NoError(t, err)
require.Equal(t, []int64{7}, repo.deletedIDs)
require.Equal(t, []int64{7}, apiKeyRepo.listByUserIDCalls)
require.Equal(t, []int64{11, 12}, apiKeyRepo.deletedIDs)
require.ElementsMatch(t, []string{"sk-user-1", "sk-user-2"}, invalidator.keys)
require.Equal(t, []int64{7}, invalidator.userIDs)
}
func TestAdminService_DeleteUser_NotFound(t *testing.T) {
repo := &userRepoStub{getErr: ErrUserNotFound}
svc := &adminServiceImpl{userRepo: repo}
err := svc.DeleteUser(context.Background(), 404)
require.ErrorIs(t, err, ErrUserNotFound)
require.Empty(t, repo.deletedIDs)
}
func TestAdminService_DeleteUser_AdminGuard(t *testing.T) {
repo := &userRepoStub{user: &User{ID: 1, Role: RoleAdmin}}
svc := &adminServiceImpl{userRepo: repo}
err := svc.DeleteUser(context.Background(), 1)
require.Error(t, err)
require.ErrorContains(t, err, "cannot delete admin user")
require.Empty(t, repo.deletedIDs)
}
func TestAdminService_DeleteUser_DeleteError(t *testing.T) {
deleteErr := errors.New("delete failed")
repo := &userRepoStub{
user: &User{ID: 9, Role: RoleUser},
deleteErr: deleteErr,
}
svc := &adminServiceImpl{userRepo: repo}
err := svc.DeleteUser(context.Background(), 9)
require.ErrorIs(t, err, deleteErr)
require.Equal(t, []int64{9}, repo.deletedIDs)
}
func TestAdminService_DeleteGroup_Success_WithCacheInvalidation(t *testing.T) {
cache := newBillingCacheStub(2)
repo := &groupRepoStub{affectedUserIDs: []int64{11, 12}}
svc := &adminServiceImpl{
groupRepo: repo,
billingCacheService: &BillingCacheService{cache: cache},
}
err := svc.DeleteGroup(context.Background(), 5)
require.NoError(t, err)
require.Equal(t, []int64{5}, repo.deleteCalls)
calls := waitForInvalidations(t, cache.invalidations, 2)
require.ElementsMatch(t, []subscriptionInvalidateCall{
{userID: 11, groupID: 5},
{userID: 12, groupID: 5},
}, calls)
}
func TestAdminService_DeleteGroup_InvalidatesAuthCacheForBoundKeys(t *testing.T) {
repo := &groupRepoStub{}
apiKeyRepo := &deleteGroupAPIKeyRepoStub{keys: []string{"k1", "k2"}}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
groupRepo: repo,
apiKeyRepo: apiKeyRepo,
authCacheInvalidator: invalidator,
}
err := svc.DeleteGroup(context.Background(), 5)
require.NoError(t, err)
require.Equal(t, []int64{5}, repo.deleteCalls)
require.Equal(t, []int64{5}, apiKeyRepo.listGroupIDs)
require.Equal(t, []string{"k1", "k2"}, invalidator.keys)
}
func TestAdminService_DeleteGroup_NotFound(t *testing.T) {
repo := &groupRepoStub{deleteErr: ErrGroupNotFound}
svc := &adminServiceImpl{groupRepo: repo}
err := svc.DeleteGroup(context.Background(), 99)
require.ErrorIs(t, err, ErrGroupNotFound)
}
func TestAdminService_DeleteGroup_Error(t *testing.T) {
deleteErr := errors.New("delete failed")
repo := &groupRepoStub{deleteErr: deleteErr}
svc := &adminServiceImpl{groupRepo: repo}
err := svc.DeleteGroup(context.Background(), 42)
require.ErrorIs(t, err, deleteErr)
}
func TestAdminService_DeleteProxy_Success(t *testing.T) {
repo := &proxyRepoStub{}
svc := &adminServiceImpl{proxyRepo: repo}
err := svc.DeleteProxy(context.Background(), 7)
require.NoError(t, err)
require.Equal(t, []int64{7}, repo.deletedIDs)
}
func TestAdminService_DeleteProxy_Idempotent(t *testing.T) {
repo := &proxyRepoStub{}
svc := &adminServiceImpl{proxyRepo: repo}
err := svc.DeleteProxy(context.Background(), 404)
require.NoError(t, err)
require.Equal(t, []int64{404}, repo.deletedIDs)
}
func TestAdminService_DeleteProxy_InUse(t *testing.T) {
repo := &proxyRepoStub{accountCount: 2}
svc := &adminServiceImpl{proxyRepo: repo}
err := svc.DeleteProxy(context.Background(), 77)
require.ErrorIs(t, err, ErrProxyInUse)
require.Empty(t, repo.deletedIDs)
}
func TestAdminService_DeleteProxy_Error(t *testing.T) {
deleteErr := errors.New("delete failed")
repo := &proxyRepoStub{deleteErr: deleteErr}
svc := &adminServiceImpl{proxyRepo: repo}
err := svc.DeleteProxy(context.Background(), 33)
require.ErrorIs(t, err, deleteErr)
}
func TestAdminService_DeleteRedeemCode_Success(t *testing.T) {
repo := &redeemRepoStub{}
svc := &adminServiceImpl{redeemCodeRepo: repo}
err := svc.DeleteRedeemCode(context.Background(), 10)
require.NoError(t, err)
require.Equal(t, []int64{10}, repo.deletedIDs)
}
func TestAdminService_DeleteRedeemCode_Idempotent(t *testing.T) {
repo := &redeemRepoStub{}
svc := &adminServiceImpl{redeemCodeRepo: repo}
err := svc.DeleteRedeemCode(context.Background(), 999)
require.NoError(t, err)
require.Equal(t, []int64{999}, repo.deletedIDs)
}
func TestAdminService_DeleteRedeemCode_Error(t *testing.T) {
deleteErr := errors.New("delete failed")
repo := &redeemRepoStub{deleteErrByID: map[int64]error{1: deleteErr}}
svc := &adminServiceImpl{redeemCodeRepo: repo}
err := svc.DeleteRedeemCode(context.Background(), 1)
require.ErrorIs(t, err, deleteErr)
require.Equal(t, []int64{1}, repo.deletedIDs)
}
func TestAdminService_BatchDeleteRedeemCodes_Success(t *testing.T) {
repo := &redeemRepoStub{}
svc := &adminServiceImpl{redeemCodeRepo: repo}
deleted, err := svc.BatchDeleteRedeemCodes(context.Background(), []int64{1, 2, 3})
require.NoError(t, err)
require.Equal(t, int64(3), deleted)
require.Equal(t, []int64{1, 2, 3}, repo.deletedIDs)
}
func TestAdminService_BatchDeleteRedeemCodes_PartialFailures(t *testing.T) {
repo := &redeemRepoStub{
deleteErrByID: map[int64]error{
2: errors.New("db error"),
},
}
svc := &adminServiceImpl{redeemCodeRepo: repo}
deleted, err := svc.BatchDeleteRedeemCodes(context.Background(), []int64{1, 2, 3})
require.NoError(t, err)
require.Equal(t, int64(2), deleted)
require.Equal(t, []int64{1, 2, 3}, repo.deletedIDs)
}
@@ -0,0 +1,326 @@
//go:build unit
package service
import (
"context"
"errors"
"net/http"
"strings"
"testing"
"time"
"unicode/utf8"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
)
type duplicateAccountRepoStub struct {
*sparkShadowRepoStub
atomicCreateErr error
accountGroupsOf map[int64][]AccountGroup
}
func newDuplicateAccountRepoStub() *duplicateAccountRepoStub {
return &duplicateAccountRepoStub{
sparkShadowRepoStub: newSparkShadowRepoStub(),
accountGroupsOf: make(map[int64][]AccountGroup),
}
}
func (s *duplicateAccountRepoStub) CreateWithAccountGroups(ctx context.Context, account *Account, groups []AccountGroup) error {
if s.atomicCreateErr != nil {
return s.atomicCreateErr
}
groupIDs := make([]int64, 0, len(groups))
for _, group := range groups {
groupIDs = append(groupIDs, group.GroupID)
}
account.GroupIDs = groupIDs
if err := s.Create(ctx, account); err != nil {
return err
}
clonedGroups := make([]AccountGroup, len(groups))
copy(clonedGroups, groups)
for i := range clonedGroups {
clonedGroups[i].AccountID = account.ID
}
account.AccountGroups = clonedGroups
s.accountGroupsOf[account.ID] = clonedGroups
if len(groupIDs) > 0 {
s.groupsOf[account.ID] = append([]int64(nil), groupIDs...)
}
stored := *account
s.accounts[account.ID] = &stored
s.mockAccountRepoForGemini.accountsByID[account.ID] = &stored
return nil
}
func (s *duplicateAccountRepoStub) FindByExtraField(_ context.Context, key string, value any) ([]Account, error) {
wanted, ok := value.(string)
if !ok {
return nil, nil
}
var matches []Account
for _, account := range s.accounts {
if actual, ok := account.Extra[key].(string); ok && actual == wanted {
matches = append(matches, *account)
}
}
return matches, nil
}
func TestDuplicateAccountCopiesConfigurationAndResetsRuntimeState(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
notes := "keep this note"
proxyID := int64(17)
originalProxyID := int64(11)
rateMultiplier := 1.25
loadFactor := 9
expiresAt := time.Date(2027, time.March, 4, 5, 6, 7, 0, time.UTC)
rateLimitedAt := time.Now().Add(-time.Minute)
rateLimitResetAt := time.Now().Add(time.Hour)
overloadUntil := time.Now().Add(2 * time.Hour)
tempUnschedulableUntil := time.Now().Add(3 * time.Hour)
sessionWindowStart := time.Now().Add(-2 * time.Hour)
sessionWindowEnd := time.Now().Add(2 * time.Hour)
source := &Account{
Name: "primary",
Notes: &notes,
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
ProxyID: &proxyID,
ProxyFallbackOriginID: &originalProxyID,
Concurrency: 6,
Priority: 40,
RateMultiplier: &rateMultiplier,
LoadFactor: &loadFactor,
Status: StatusError,
Schedulable: true,
ErrorMessage: "upstream unavailable",
ExpiresAt: &expiresAt,
AutoPauseOnExpired: false,
Credentials: map[string]any{
"api_key": "secret",
"nested": map[string]any{"token": "source-token"},
},
Extra: map[string]any{
"config": map[string]any{"region": "us-east-1"},
"items": []any{map[string]any{"enabled": true}},
"quota_limit": 1000,
"quota_used": 450,
"quota_daily_used": 25,
"quota_daily_start": "2026-07-15T00:00:00Z",
"model_rate_limits": map[string]any{"gpt-5": "2099-01-01T00:00:00Z"},
"codex_5h_used_percent": 80,
"codex_cli_only": true,
"grok_usage_snapshot": map[string]any{"status_code": 429},
"openai_responses_supported": false,
"openai_compact_checked_at": "2026-07-15T00:00:00Z",
"session_window_utilization": 0.8,
"passive_usage_sampled_at": "2026-07-15T00:00:00Z",
"antigravity_force_token_refresh": true,
"antigravity_credits_overages": map[string]any{"enabled": true},
"crs_account_id": "remote-42",
"crs_kind": "openai-api-key",
"crs_synced_at": "2026-07-15T00:00:00Z",
},
GroupIDs: []int64{7, 3},
AccountGroups: []AccountGroup{{GroupID: 7, Priority: 50}, {GroupID: 3, Priority: 7}},
RateLimitedAt: &rateLimitedAt,
RateLimitResetAt: &rateLimitResetAt,
OverloadUntil: &overloadUntil,
TempUnschedulableUntil: &tempUnschedulableUntil,
TempUnschedulableReason: "maintenance",
SessionWindowStart: &sessionWindowStart,
SessionWindowEnd: &sessionWindowEnd,
SessionWindowStatus: "active",
}
source.Extra[UpstreamBillingProbeEnabledExtraKey] = true
source.Extra[UpstreamBillingRateSyncEnabledExtraKey] = true
source.Extra[UpstreamBillingProbeExtraKey] = map[string]any{"status": "ok"}
require.NoError(t, repo.Create(ctx, source))
duplicate, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
require.NoError(t, err)
require.NotEqual(t, source.ID, duplicate.ID)
require.Equal(t, "primary (Copy)", duplicate.Name)
require.Equal(t, source.Platform, duplicate.Platform)
require.Equal(t, source.Type, duplicate.Type)
require.Equal(t, source.Concurrency, duplicate.Concurrency)
require.Equal(t, source.Priority, duplicate.Priority)
require.Equal(t, source.AutoPauseOnExpired, duplicate.AutoPauseOnExpired)
require.Equal(t, source.GroupIDs, duplicate.GroupIDs)
require.Equal(t, source.Credentials, duplicate.Credentials)
require.Equal(t, map[string]any{
"config": map[string]any{"region": "us-east-1"},
"items": []any{map[string]any{"enabled": true}},
"quota_limit": float64(1000),
"codex_cli_only": true,
}, duplicate.Extra)
require.NotContains(t, duplicate.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotNil(t, duplicate.ExpiresAt)
require.True(t, source.ExpiresAt.Equal(*duplicate.ExpiresAt))
require.Equal(t, source.Notes, duplicate.Notes)
require.Equal(t, source.ProxyFallbackOriginID, duplicate.ProxyID)
require.Equal(t, source.RateMultiplier, duplicate.RateMultiplier)
require.Equal(t, source.LoadFactor, duplicate.LoadFactor)
require.Equal(t, source.GroupIDs, repo.groupsOf[duplicate.ID])
require.Equal(t, []AccountGroup{
{AccountID: duplicate.ID, GroupID: 7, Priority: 50},
{AccountID: duplicate.ID, GroupID: 3, Priority: 7},
}, repo.accountGroupsOf[duplicate.ID])
require.Equal(t, StatusActive, duplicate.Status)
require.False(t, duplicate.Schedulable)
require.Empty(t, duplicate.ErrorMessage)
require.Nil(t, duplicate.LastUsedAt)
require.Nil(t, duplicate.RateLimitedAt)
require.Nil(t, duplicate.RateLimitResetAt)
require.Nil(t, duplicate.OverloadUntil)
require.Nil(t, duplicate.TempUnschedulableUntil)
require.Empty(t, duplicate.TempUnschedulableReason)
require.Nil(t, duplicate.SessionWindowStart)
require.Nil(t, duplicate.SessionWindowEnd)
require.Empty(t, duplicate.SessionWindowStatus)
duplicate.Credentials["nested"].(map[string]any)["token"] = "changed"
duplicate.Extra["config"].(map[string]any)["region"] = "changed"
duplicate.Extra["items"].([]any)[0].(map[string]any)["enabled"] = false
storedSource, getErr := repo.GetByID(ctx, source.ID)
require.NoError(t, getErr)
require.Equal(t, "source-token", storedSource.Credentials["nested"].(map[string]any)["token"])
require.Equal(t, "us-east-1", storedSource.Extra["config"].(map[string]any)["region"])
require.Equal(t, true, storedSource.Extra["items"].([]any)[0].(map[string]any)["enabled"])
require.Equal(t, "remote-42", storedSource.Extra["crs_account_id"])
}
func TestDuplicateAccountRejectsCredentialShadow(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
parentID := int64(99)
shadow := &Account{
Name: "shadow",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
ParentAccountID: &parentID,
QuotaDimension: QuotaDimensionSpark,
}
require.NoError(t, repo.Create(ctx, shadow))
_, err := svc.DuplicateAccount(ctx, shadow.ID, "admin:1", "")
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Equal(t, "ACCOUNT_DUPLICATE_SHADOW_UNSUPPORTED", infraerrors.Reason(err))
require.Len(t, repo.accounts, 1)
}
func TestDuplicateAccountRejectsRotatingOrUnknownCredentialTypes(t *testing.T) {
for _, accountType := range []string{AccountTypeOAuth, AccountTypeSetupToken, "legacy-cookie"} {
t.Run(accountType, func(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
source := &Account{
Name: "rotating-credential-account",
Platform: PlatformOpenAI,
Type: accountType,
Credentials: map[string]any{"refresh_token": "shared-token"},
}
require.NoError(t, repo.Create(ctx, source))
_, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Equal(t, "ACCOUNT_DUPLICATE_CREDENTIAL_TYPE_UNSUPPORTED", infraerrors.Reason(err))
require.Len(t, repo.accounts, 1)
})
}
}
func TestDuplicateAccountPreservesUngroupedState(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
source := &Account{
Name: "ungrouped",
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "secret"},
GroupIDs: nil,
}
require.NoError(t, repo.Create(ctx, source))
duplicate, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
require.NoError(t, err)
require.Empty(t, duplicate.GroupIDs)
require.NotContains(t, repo.groupsOf, duplicate.ID)
}
func TestDuplicateAccountAtomicCreateFailureLeavesNoOrphan(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
source := &Account{
Name: "source",
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "secret"},
GroupIDs: []int64{7},
AccountGroups: []AccountGroup{{GroupID: 7, Priority: 25}},
}
require.NoError(t, repo.Create(ctx, source))
repo.atomicCreateErr = errors.New("group binding failed")
_, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
require.ErrorContains(t, err, "group binding failed")
require.Len(t, repo.accounts, 1)
}
func TestDuplicateAccountNamePreservesSuffixWithinSchemaLimit(t *testing.T) {
name := duplicateAccountName(strings.Repeat("界", 100))
require.Equal(t, 100, utf8.RuneCountInString(name))
require.True(t, strings.HasSuffix(name, " (Copy)"))
}
func TestDuplicateAccountReturnsExistingCopyForSameOperationKey(t *testing.T) {
ctx := context.Background()
repo := newDuplicateAccountRepoStub()
svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
source := &Account{
Name: "source",
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "secret"},
}
require.NoError(t, repo.Create(ctx, source))
first, err := svc.DuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
require.NoError(t, err)
second, err := svc.DuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
require.NoError(t, err)
recovered, err := svc.RecoverDuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
require.NoError(t, err)
otherAdminRecovery, err := svc.RecoverDuplicateAccount(ctx, source.ID, "admin:8", "stable-operation-key")
require.NoError(t, err)
otherAdminCopy, err := svc.DuplicateAccount(ctx, source.ID, "admin:8", "stable-operation-key")
require.NoError(t, err)
require.Equal(t, first.ID, second.ID)
require.Equal(t, first.ID, recovered.ID)
require.Nil(t, otherAdminRecovery, "durable recovery identity must remain scoped to the initiating admin")
require.NotEqual(t, first.ID, otherAdminCopy.ID)
require.Len(t, repo.accounts, 3)
require.NotEmpty(t, first.Extra[duplicateAccountOperationIDExtraKey])
}
@@ -0,0 +1,216 @@
//go:build unit
package service
import (
"context"
"fmt"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type ensureEmailCall struct {
userID int64
email string
}
type replaceEmailCall struct {
userID int64
oldEmail string
newEmail string
}
type emailSyncRepoStub struct {
user *User
nextID int64
updateCalls int
created []*User
updated []*User
ensureCalls []ensureEmailCall
replaceCalls []replaceEmailCall
ensureErr error
replaceErr error
}
func (s *emailSyncRepoStub) CreateWithEmailAliasGuard(ctx context.Context, user *User) error {
return s.Create(ctx, user)
}
func (s *emailSyncRepoStub) Create(_ context.Context, user *User) error {
if s.nextID != 0 && user.ID == 0 {
user.ID = s.nextID
}
s.created = append(s.created, user)
s.user = user
return nil
}
func (s *emailSyncRepoStub) GetByID(_ context.Context, _ int64) (*User, error) {
if s.user == nil {
return nil, ErrUserNotFound
}
cloned := *s.user
return &cloned, nil
}
func (s *emailSyncRepoStub) GetByEmail(_ context.Context, _ string) (*User, error) {
return nil, ErrUserNotFound
}
func (s *emailSyncRepoStub) GetFirstAdmin(context.Context) (*User, error) {
return nil, fmt.Errorf("unexpected GetFirstAdmin call")
}
func (s *emailSyncRepoStub) Update(_ context.Context, user *User, _ UserUpdateFields) error {
s.updateCalls++
s.updated = append(s.updated, user)
s.user = user
return nil
}
func (s *emailSyncRepoStub) Delete(context.Context, int64) error { return nil }
func (s *emailSyncRepoStub) GetUserAvatar(context.Context, int64) (*UserAvatar, error) {
return nil, fmt.Errorf("unexpected GetUserAvatar call")
}
func (s *emailSyncRepoStub) UpsertUserAvatar(context.Context, int64, UpsertUserAvatarInput) (*UserAvatar, error) {
return nil, fmt.Errorf("unexpected UpsertUserAvatar call")
}
func (s *emailSyncRepoStub) DeleteUserAvatar(context.Context, int64) error {
return fmt.Errorf("unexpected DeleteUserAvatar call")
}
func (s *emailSyncRepoStub) List(context.Context, pagination.PaginationParams) ([]User, *pagination.PaginationResult, error) {
return nil, nil, fmt.Errorf("unexpected List call")
}
func (s *emailSyncRepoStub) ListWithFilters(context.Context, pagination.PaginationParams, UserListFilters) ([]User, *pagination.PaginationResult, error) {
return nil, nil, fmt.Errorf("unexpected ListWithFilters call")
}
func (s *emailSyncRepoStub) GetLatestUsedAtByUserIDs(context.Context, []int64) (map[int64]*time.Time, error) {
return map[int64]*time.Time{}, nil
}
func (s *emailSyncRepoStub) GetLatestUsedAtByUserID(context.Context, int64) (*time.Time, error) {
return nil, nil
}
func (s *emailSyncRepoStub) UpdateUserLastActiveAt(context.Context, int64, time.Time) error {
return nil
}
func (s *emailSyncRepoStub) UpdateBalance(context.Context, int64, float64) error { return nil }
func (s *emailSyncRepoStub) DeductBalance(context.Context, int64, float64) error { return nil }
func (s *emailSyncRepoStub) UpdateConcurrency(context.Context, int64, int) error { return nil }
func (s *emailSyncRepoStub) ExistsByEmail(context.Context, string) (bool, error) { return false, nil }
func (s *emailSyncRepoStub) ExistsByEmailAlias(context.Context, string) (bool, error) {
return false, nil
}
func (s *emailSyncRepoStub) AdjustBalance(ctx context.Context, id int64, delta float64) (BalanceChange, error) {
panic("unexpected AdjustBalance call")
}
func (s *emailSyncRepoStub) SetBalance(ctx context.Context, id int64, value float64) (BalanceChange, error) {
panic("unexpected SetBalance call")
}
func (s *emailSyncRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
return 0, nil
}
func (s *emailSyncRepoStub) BatchSetConcurrency(context.Context, []int64, int) (int, error) {
return 0, nil
}
func (s *emailSyncRepoStub) BatchAddConcurrency(context.Context, []int64, int) (int, error) {
return 0, nil
}
func (s *emailSyncRepoStub) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
return 0, nil
}
func (s *emailSyncRepoStub) AddGroupToAllowedGroups(context.Context, int64, int64) error { return nil }
func (s *emailSyncRepoStub) RemoveGroupFromUserAllowedGroups(context.Context, int64, int64) error {
return nil
}
func (s *emailSyncRepoStub) ListUserAuthIdentities(context.Context, int64) ([]UserAuthIdentityRecord, error) {
return nil, nil
}
func (s *emailSyncRepoStub) UnbindUserAuthProvider(context.Context, int64, string) error { return nil }
func (s *emailSyncRepoStub) UpdateTotpSecret(context.Context, int64, *string) error { return nil }
func (s *emailSyncRepoStub) EnableTotp(context.Context, int64) error { return nil }
func (s *emailSyncRepoStub) DisableTotp(context.Context, int64) error { return nil }
func (s *emailSyncRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) {
return s.GetByID(ctx, id)
}
func (s *emailSyncRepoStub) EnsureEmailAuthIdentity(_ context.Context, userID int64, email string) error {
s.ensureCalls = append(s.ensureCalls, ensureEmailCall{userID: userID, email: email})
return s.ensureErr
}
func (s *emailSyncRepoStub) ReplaceEmailAuthIdentity(_ context.Context, userID int64, oldEmail, newEmail string) error {
s.replaceCalls = append(s.replaceCalls, replaceEmailCall{
userID: userID,
oldEmail: oldEmail,
newEmail: newEmail,
})
return s.replaceErr
}
func TestAdminService_CreateUser_DoesNotReturnPartialSuccessFromEmailIdentityResync(t *testing.T) {
repo := &emailSyncRepoStub{
nextID: 55,
ensureErr: fmt.Errorf("unexpected email resync"),
}
svc := &adminServiceImpl{userRepo: repo}
user, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "admin-created@example.com",
Password: "strong-pass",
})
require.NoError(t, err)
require.NotNil(t, user)
require.Equal(t, int64(55), user.ID)
require.Empty(t, repo.ensureCalls)
require.Empty(t, repo.replaceCalls)
}
func TestAdminService_UpdateUser_DoesNotReturnPartialSuccessFromEmailIdentityResync(t *testing.T) {
repo := &emailSyncRepoStub{
user: &User{
ID: 91,
Email: "before@example.com",
Role: RoleUser,
Status: StatusActive,
Concurrency: 3,
},
replaceErr: fmt.Errorf("unexpected email resync"),
}
svc := &adminServiceImpl{userRepo: repo}
updated, err := svc.UpdateUser(context.Background(), 91, &UpdateUserInput{
Email: "after@example.com",
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, "after@example.com", updated.Email)
require.Empty(t, repo.replaceCalls)
require.Empty(t, repo.ensureCalls)
}
@@ -0,0 +1,22 @@
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestAdminService_GetUserIncludeDeleted(t *testing.T) {
ts := time.Date(2026, 5, 28, 0, 0, 0, 0, time.UTC)
repo := &userRepoStub{user: &User{ID: 7, Email: "del@test.com", DeletedAt: &ts}}
svc := &adminServiceImpl{userRepo: repo}
got, err := svc.GetUserIncludeDeleted(context.Background(), 7)
require.NoError(t, err)
require.Equal(t, int64(7), got.ID)
require.NotNil(t, got.DeletedAt)
}
@@ -0,0 +1,225 @@
//go:build unit
package service
import (
"context"
"errors"
"net/http"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
)
// userGroupRateRepoStubForGroupRate implements UserGroupRateRepository for group rate tests.
type userGroupRateRepoStubForGroupRate struct {
getByGroupIDData map[int64][]UserGroupRateEntry
getByGroupIDErr error
deletedGroupIDs []int64
deleteByGroupErr error
syncedGroupID int64
syncedEntries []GroupRateMultiplierInput
syncGroupErr error
rpmSyncedGroupID int64
rpmSyncedEntries []GroupRPMOverrideInput
rpmSyncErr error
}
func (s *userGroupRateRepoStubForGroupRate) GetByUserID(_ context.Context, _ int64) (map[int64]float64, error) {
panic("unexpected GetByUserID call")
}
func (s *userGroupRateRepoStubForGroupRate) GetByUserAndGroup(_ context.Context, _, _ int64) (*float64, error) {
panic("unexpected GetByUserAndGroup call")
}
func (s *userGroupRateRepoStubForGroupRate) GetRPMOverrideByUserAndGroup(_ context.Context, _, _ int64) (*int, error) {
panic("unexpected GetRPMOverrideByUserAndGroup call")
}
func (s *userGroupRateRepoStubForGroupRate) GetByGroupID(_ context.Context, groupID int64) ([]UserGroupRateEntry, error) {
if s.getByGroupIDErr != nil {
return nil, s.getByGroupIDErr
}
return s.getByGroupIDData[groupID], nil
}
func (s *userGroupRateRepoStubForGroupRate) SyncUserGroupRates(_ context.Context, _ int64, _ map[int64]*float64) error {
panic("unexpected SyncUserGroupRates call")
}
func (s *userGroupRateRepoStubForGroupRate) SyncGroupRateMultipliers(_ context.Context, groupID int64, entries []GroupRateMultiplierInput) error {
s.syncedGroupID = groupID
s.syncedEntries = entries
return s.syncGroupErr
}
func (s *userGroupRateRepoStubForGroupRate) SyncGroupRPMOverrides(_ context.Context, groupID int64, entries []GroupRPMOverrideInput) error {
s.rpmSyncedGroupID = groupID
s.rpmSyncedEntries = entries
return s.rpmSyncErr
}
func (s *userGroupRateRepoStubForGroupRate) ClearGroupRPMOverrides(_ context.Context, _ int64) error {
panic("unexpected ClearGroupRPMOverrides call")
}
func (s *userGroupRateRepoStubForGroupRate) DeleteByGroupID(_ context.Context, groupID int64) error {
s.deletedGroupIDs = append(s.deletedGroupIDs, groupID)
return s.deleteByGroupErr
}
func (s *userGroupRateRepoStubForGroupRate) DeleteByUserID(_ context.Context, _ int64) error {
panic("unexpected DeleteByUserID call")
}
func TestAdminService_GetGroupRateMultipliers(t *testing.T) {
t.Run("returns entries for group", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{
getByGroupIDData: map[int64][]UserGroupRateEntry{
10: {
{UserID: 1, UserName: "alice", UserEmail: "alice@test.com", RateMultiplier: ptrFloat(1.5)},
{UserID: 2, UserName: "bob", UserEmail: "bob@test.com", RateMultiplier: ptrFloat(0.8)},
},
},
}
svc := &adminServiceImpl{userGroupRateRepo: repo}
entries, err := svc.GetGroupRateMultipliers(context.Background(), 10)
require.NoError(t, err)
require.Len(t, entries, 2)
require.Equal(t, int64(1), entries[0].UserID)
require.Equal(t, "alice", entries[0].UserName)
require.NotNil(t, entries[0].RateMultiplier)
require.Equal(t, 1.5, *entries[0].RateMultiplier)
require.Equal(t, int64(2), entries[1].UserID)
require.NotNil(t, entries[1].RateMultiplier)
require.Equal(t, 0.8, *entries[1].RateMultiplier)
})
t.Run("returns nil when repo is nil", func(t *testing.T) {
svc := &adminServiceImpl{userGroupRateRepo: nil}
entries, err := svc.GetGroupRateMultipliers(context.Background(), 10)
require.NoError(t, err)
require.Nil(t, entries)
})
t.Run("returns empty slice for group with no entries", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{
getByGroupIDData: map[int64][]UserGroupRateEntry{},
}
svc := &adminServiceImpl{userGroupRateRepo: repo}
entries, err := svc.GetGroupRateMultipliers(context.Background(), 99)
require.NoError(t, err)
require.Nil(t, entries)
})
t.Run("propagates repo error", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{
getByGroupIDErr: errors.New("db error"),
}
svc := &adminServiceImpl{userGroupRateRepo: repo}
_, err := svc.GetGroupRateMultipliers(context.Background(), 10)
require.Error(t, err)
require.Contains(t, err.Error(), "db error")
})
}
func TestAdminService_ClearGroupRateMultipliers(t *testing.T) {
t.Run("deletes by group ID", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{}
svc := &adminServiceImpl{userGroupRateRepo: repo}
err := svc.ClearGroupRateMultipliers(context.Background(), 42)
require.NoError(t, err)
require.Equal(t, []int64{42}, repo.deletedGroupIDs)
})
t.Run("returns nil when repo is nil", func(t *testing.T) {
svc := &adminServiceImpl{userGroupRateRepo: nil}
err := svc.ClearGroupRateMultipliers(context.Background(), 42)
require.NoError(t, err)
})
t.Run("propagates repo error", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{
deleteByGroupErr: errors.New("delete failed"),
}
svc := &adminServiceImpl{userGroupRateRepo: repo}
err := svc.ClearGroupRateMultipliers(context.Background(), 42)
require.Error(t, err)
require.Contains(t, err.Error(), "delete failed")
})
}
func TestAdminService_BatchSetGroupRateMultipliers(t *testing.T) {
t.Run("syncs entries to repo", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{}
svc := &adminServiceImpl{userGroupRateRepo: repo}
entries := []GroupRateMultiplierInput{
{UserID: 1, RateMultiplier: 1.5},
{UserID: 2, RateMultiplier: 0.8},
}
err := svc.BatchSetGroupRateMultipliers(context.Background(), 10, entries)
require.NoError(t, err)
require.Equal(t, int64(10), repo.syncedGroupID)
require.Equal(t, entries, repo.syncedEntries)
})
t.Run("returns nil when repo is nil", func(t *testing.T) {
svc := &adminServiceImpl{userGroupRateRepo: nil}
err := svc.BatchSetGroupRateMultipliers(context.Background(), 10, nil)
require.NoError(t, err)
})
t.Run("propagates repo error", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{
syncGroupErr: errors.New("sync failed"),
}
svc := &adminServiceImpl{userGroupRateRepo: repo}
err := svc.BatchSetGroupRateMultipliers(context.Background(), 10, []GroupRateMultiplierInput{
{UserID: 1, RateMultiplier: 1.0},
})
require.Error(t, err)
require.Contains(t, err.Error(), "sync failed")
})
}
func TestAdminService_BatchSetGroupRPMOverrides(t *testing.T) {
t.Run("syncs entries to repo", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{}
svc := &adminServiceImpl{userGroupRateRepo: repo}
override := 20
entries := []GroupRPMOverrideInput{{UserID: 2, RPMOverride: &override}}
err := svc.BatchSetGroupRPMOverrides(context.Background(), 10, entries)
require.NoError(t, err)
require.Equal(t, int64(10), repo.rpmSyncedGroupID)
require.Equal(t, entries, repo.rpmSyncedEntries)
})
t.Run("rejects negative override as bad request", func(t *testing.T) {
repo := &userGroupRateRepoStubForGroupRate{}
svc := &adminServiceImpl{userGroupRateRepo: repo}
negative := -1
err := svc.BatchSetGroupRPMOverrides(context.Background(), 10, []GroupRPMOverrideInput{
{UserID: 2, RPMOverride: &negative},
})
require.Error(t, err)
require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
require.Zero(t, repo.rpmSyncedGroupID)
})
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,185 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type userRepoStubForListUsers struct {
userRepoStub
users []User
err error
listWithFiltersParams pagination.PaginationParams
lastUsedByUserID map[int64]*time.Time
lastUsedErr error
}
func (s *userRepoStubForListUsers) ListWithFilters(_ context.Context, params pagination.PaginationParams, _ UserListFilters) ([]User, *pagination.PaginationResult, error) {
s.listWithFiltersParams = params
if s.err != nil {
return nil, nil, s.err
}
out := make([]User, len(s.users))
copy(out, s.users)
return out, &pagination.PaginationResult{
Total: int64(len(out)),
Page: params.Page,
PageSize: params.PageSize,
}, nil
}
func (s *userRepoStubForListUsers) GetLatestUsedAtByUserIDs(_ context.Context, userIDs []int64) (map[int64]*time.Time, error) {
if s.lastUsedErr != nil {
return nil, s.lastUsedErr
}
result := make(map[int64]*time.Time, len(userIDs))
for _, userID := range userIDs {
if ts, ok := s.lastUsedByUserID[userID]; ok {
result[userID] = ts
}
}
return result, nil
}
func (s *userRepoStubForListUsers) GetLatestUsedAtByUserID(_ context.Context, userID int64) (*time.Time, error) {
if s.lastUsedErr != nil {
return nil, s.lastUsedErr
}
return s.lastUsedByUserID[userID], nil
}
type userGroupRateRepoStubForListUsers struct {
batchCalls int
singleCall []int64
batchErr error
batchData map[int64]map[int64]float64
singleErr map[int64]error
singleData map[int64]map[int64]float64
}
func (s *userGroupRateRepoStubForListUsers) GetByUserIDs(_ context.Context, _ []int64) (map[int64]map[int64]float64, error) {
s.batchCalls++
if s.batchErr != nil {
return nil, s.batchErr
}
return s.batchData, nil
}
func (s *userGroupRateRepoStubForListUsers) GetByUserID(_ context.Context, userID int64) (map[int64]float64, error) {
s.singleCall = append(s.singleCall, userID)
if err, ok := s.singleErr[userID]; ok {
return nil, err
}
if rates, ok := s.singleData[userID]; ok {
return rates, nil
}
return map[int64]float64{}, nil
}
func (s *userGroupRateRepoStubForListUsers) GetByUserAndGroup(_ context.Context, userID, groupID int64) (*float64, error) {
panic("unexpected GetByUserAndGroup call")
}
func (s *userGroupRateRepoStubForListUsers) GetRPMOverrideByUserAndGroup(_ context.Context, _, _ int64) (*int, error) {
panic("unexpected GetRPMOverrideByUserAndGroup call")
}
func (s *userGroupRateRepoStubForListUsers) SyncUserGroupRates(_ context.Context, userID int64, rates map[int64]*float64) error {
panic("unexpected SyncUserGroupRates call")
}
func (s *userGroupRateRepoStubForListUsers) GetByGroupID(_ context.Context, _ int64) ([]UserGroupRateEntry, error) {
panic("unexpected GetByGroupID call")
}
func (s *userGroupRateRepoStubForListUsers) SyncGroupRateMultipliers(_ context.Context, _ int64, _ []GroupRateMultiplierInput) error {
panic("unexpected SyncGroupRateMultipliers call")
}
func (s *userGroupRateRepoStubForListUsers) SyncGroupRPMOverrides(_ context.Context, _ int64, _ []GroupRPMOverrideInput) error {
panic("unexpected SyncGroupRPMOverrides call")
}
func (s *userGroupRateRepoStubForListUsers) ClearGroupRPMOverrides(_ context.Context, _ int64) error {
panic("unexpected ClearGroupRPMOverrides call")
}
func (s *userGroupRateRepoStubForListUsers) DeleteByGroupID(_ context.Context, _ int64) error {
panic("unexpected DeleteByGroupID call")
}
func (s *userGroupRateRepoStubForListUsers) DeleteByUserID(_ context.Context, userID int64) error {
panic("unexpected DeleteByUserID call")
}
func TestAdminService_ListUsers_BatchRateFallbackToSingle(t *testing.T) {
userRepo := &userRepoStubForListUsers{
users: []User{
{ID: 101, Username: "u1"},
{ID: 202, Username: "u2"},
},
}
rateRepo := &userGroupRateRepoStubForListUsers{
batchErr: errors.New("batch unavailable"),
singleData: map[int64]map[int64]float64{
101: {11: 1.1},
202: {22: 2.2},
},
}
svc := &adminServiceImpl{
userRepo: userRepo,
userGroupRateRepo: rateRepo,
}
users, total, err := svc.ListUsers(context.Background(), 1, 20, UserListFilters{}, "", "")
require.NoError(t, err)
require.Equal(t, int64(2), total)
require.Len(t, users, 2)
require.Equal(t, 1, rateRepo.batchCalls)
require.ElementsMatch(t, []int64{101, 202}, rateRepo.singleCall)
require.Equal(t, 1.1, users[0].GroupRates[11])
require.Equal(t, 2.2, users[1].GroupRates[22])
}
func TestAdminService_ListUsers_PassesSortParams(t *testing.T) {
userRepo := &userRepoStubForListUsers{
users: []User{{ID: 1, Email: "a@example.com"}},
}
svc := &adminServiceImpl{userRepo: userRepo}
_, _, err := svc.ListUsers(context.Background(), 2, 50, UserListFilters{}, "email", "ASC")
require.NoError(t, err)
require.Equal(t, pagination.PaginationParams{
Page: 2,
PageSize: 50,
SortBy: "email",
SortOrder: "ASC",
}, userRepo.listWithFiltersParams)
}
func TestAdminService_ListUsers_PopulatesLastUsedAt(t *testing.T) {
lastUsed := time.Now().UTC().Add(-30 * time.Minute).Truncate(time.Second)
userRepo := &userRepoStubForListUsers{
users: []User{{ID: 101, Email: "u@example.com"}},
lastUsedByUserID: map[int64]*time.Time{
101: &lastUsed,
},
}
svc := &adminServiceImpl{userRepo: userRepo}
users, total, err := svc.ListUsers(context.Background(), 1, 20, UserListFilters{}, "", "")
require.NoError(t, err)
require.Equal(t, int64(1), total)
require.Len(t, users, 1)
require.NotNil(t, users[0].LastUsedAt)
require.WithinDuration(t, lastUsed, *users[0].LastUsedAt, time.Second)
}
@@ -0,0 +1,194 @@
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type updateAccountOveragesRepoStub struct {
mockAccountRepoForGemini
account *Account
updateCalls int
}
func (r *updateAccountOveragesRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
return r.account, nil
}
func (r *updateAccountOveragesRepoStub) Update(ctx context.Context, account *Account) error {
r.updateCalls++
r.account = account
return nil
}
func TestUpdateAccount_DisableOveragesClearsAICreditsKey(t *testing.T) {
accountID := int64(101)
repo := &updateAccountOveragesRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{
"allow_overages": true,
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limited_at": "2026-03-15T00:00:00Z",
"rate_limit_reset_at": "2099-03-15T00:00:00Z",
},
creditsExhaustedKey: map[string]any{
"rate_limited_at": "2026-03-15T00:00:00Z",
"rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339),
},
},
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limited_at": "2026-03-15T00:00:00Z",
"rate_limit_reset_at": "2099-03-15T00:00:00Z",
},
creditsExhaustedKey: map[string]any{
"rate_limited_at": "2026-03-15T00:00:00Z",
"rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339),
},
},
},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.False(t, updated.IsOveragesEnabled())
// 关闭 overages 后,AICredits key 应被清除
rawLimits, ok := repo.account.Extra[modelRateLimitsKey].(map[string]any)
if ok {
_, exists := rawLimits[creditsExhaustedKey]
require.False(t, exists, "关闭 overages 时应清除 AICredits 限流 key")
}
// 普通模型限流应保留
require.True(t, ok)
_, exists := rawLimits["claude-sonnet-4-5"]
require.True(t, exists, "普通模型限流应保留")
}
func TestUpdateAccount_EnableOveragesClearsModelRateLimitsBeforePersist(t *testing.T) {
accountID := int64(102)
repo := &updateAccountOveragesRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limited_at": "2026-03-15T00:00:00Z",
"rate_limit_reset_at": "2099-03-15T00:00:00Z",
},
},
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
"mixed_scheduling": true,
"allow_overages": true,
},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.True(t, updated.IsOveragesEnabled())
_, exists := repo.account.Extra[modelRateLimitsKey]
require.False(t, exists, "开启 overages 时应在持久化前清掉旧模型限流")
}
func TestUpdateAccount_EmptyExtraPayloadCanClearQuotaLimits(t *testing.T) {
accountID := int64(103)
repo := &updateAccountOveragesRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{
"quota_limit": 100.0,
"quota_daily_limit": 10.0,
"quota_weekly_limit": 40.0,
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
// 显式空对象:语义是“清空 extra 中的可配置键”(例如关闭配额限制)
Extra: map[string]any{},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.NotNil(t, repo.account.Extra)
require.NotContains(t, repo.account.Extra, "quota_limit")
require.NotContains(t, repo.account.Extra, "quota_daily_limit")
require.NotContains(t, repo.account.Extra, "quota_weekly_limit")
require.Len(t, repo.account.Extra, 0)
}
func TestUpdateAccount_FixedWeeklyResetClearsLegacyRollingUsage(t *testing.T) {
now := time.Now().UTC()
daysSinceMonday := (int(now.Weekday()) + 6) % 7
currentWeekStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, -daysSinceMonday)
legacyRollingStart := currentWeekStart.Add(-24 * time.Hour)
accountID := int64(104)
repo := &updateAccountOveragesRepoStub{
account: &Account{
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{
"quota_weekly_limit": 40.0,
"quota_weekly_used": 12.5,
"quota_weekly_start": legacyRollingStart.Format(time.RFC3339),
},
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
"quota_weekly_limit": 40.0,
"quota_weekly_reset_mode": "fixed",
"quota_weekly_reset_day": float64(1),
"quota_weekly_reset_hour": float64(0),
"quota_reset_timezone": "UTC",
},
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.InDelta(t, 0.0, updated.GetQuotaWeeklyUsed(), 1e-9)
require.Equal(t, currentWeekStart.Format(time.RFC3339), updated.Extra["quota_weekly_start"])
require.False(t, updated.IsWeeklyQuotaPeriodExpired())
}
@@ -0,0 +1,136 @@
package service
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/require"
)
func TestFinalizeProxyQualityResult_ScoreAndGrade(t *testing.T) {
result := &ProxyQualityCheckResult{
PassedCount: 2,
WarnCount: 1,
FailedCount: 1,
ChallengeCount: 1,
}
finalizeProxyQualityResult(result)
require.Equal(t, 38, result.Score)
require.Equal(t, "F", result.Grade)
require.Contains(t, result.Summary, "通过 2 项")
require.Contains(t, result.Summary, "告警 1 项")
require.Contains(t, result.Summary, "失败 1 项")
require.Contains(t, result.Summary, "挑战 1 项")
}
func TestRunProxyQualityTarget_CloudflareChallenge(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/html")
w.Header().Set("cf-ray", "test-ray-123")
w.WriteHeader(http.StatusForbidden)
_, _ = w.Write([]byte("<!DOCTYPE html><title>Just a moment...</title><script>window._cf_chl_opt={};</script>"))
}))
defer server.Close()
target := proxyQualityTarget{
Target: "openai",
URL: server.URL,
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
},
}
item := runProxyQualityTarget(context.Background(), server.Client(), target)
require.Equal(t, "challenge", item.Status)
require.Equal(t, http.StatusForbidden, item.HTTPStatus)
require.Equal(t, "test-ray-123", item.CFRay)
}
func TestRunProxyQualityTarget_AllowedStatusPass(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"models":[]}`))
}))
defer server.Close()
target := proxyQualityTarget{
Target: "gemini",
URL: server.URL,
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusOK: {},
},
}
item := runProxyQualityTarget(context.Background(), server.Client(), target)
require.Equal(t, "pass", item.Status)
require.Equal(t, http.StatusOK, item.HTTPStatus)
}
func TestRunProxyQualityTarget_AllowedStatusPassForUnauthorized(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
}))
defer server.Close()
target := proxyQualityTarget{
Target: "openai",
URL: server.URL,
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
},
}
item := runProxyQualityTarget(context.Background(), server.Client(), target)
require.Equal(t, "pass", item.Status)
require.Equal(t, http.StatusUnauthorized, item.HTTPStatus)
require.Contains(t, item.Message, "目标可达")
}
func TestProxyQualityTargets_IncludesGrok(t *testing.T) {
var grokTarget *proxyQualityTarget
for i := range proxyQualityTargets {
if proxyQualityTargets[i].Target == "grok" {
grokTarget = &proxyQualityTargets[i]
break
}
}
require.NotNil(t, grokTarget)
require.Equal(t, "https://api.x.ai/v1/models", grokTarget.URL)
require.Equal(t, http.MethodGet, grokTarget.Method)
require.Contains(t, grokTarget.AllowedStatuses, http.StatusUnauthorized)
}
func TestRunProxyQualityTarget_GrokUnauthorizedPasses(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("unexpected method: %s", r.Method)
}
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
}))
defer server.Close()
target := proxyQualityTarget{
Target: "grok",
URL: server.URL,
Method: http.MethodGet,
AllowedStatuses: map[int]struct{}{
http.StatusUnauthorized: {},
},
}
item := runProxyQualityTarget(context.Background(), server.Client(), target)
require.Equal(t, "grok", item.Target)
require.Equal(t, "pass", item.Status)
require.Equal(t, http.StatusUnauthorized, item.HTTPStatus)
require.Contains(t, item.Message, "目标可达")
}
@@ -0,0 +1,143 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
func TestAdminService_CreateUser_WithAdminRole(t *testing.T) {
repo := &userRepoStub{nextID: 30}
svc := &adminServiceImpl{userRepo: repo}
user, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "admin@test.com",
Password: "strong-pass",
Role: RoleAdmin,
})
require.NoError(t, err)
require.Equal(t, RoleAdmin, user.Role)
}
func TestAdminService_CreateUser_DefaultsToUserRole(t *testing.T) {
repo := &userRepoStub{nextID: 31}
svc := &adminServiceImpl{userRepo: repo}
user, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "plain@test.com",
Password: "strong-pass",
})
require.NoError(t, err)
require.Equal(t, RoleUser, user.Role)
}
func TestAdminService_CreateUser_InvalidRoleRejected(t *testing.T) {
repo := &userRepoStub{nextID: 32}
svc := &adminServiceImpl{userRepo: repo}
_, err := svc.CreateUser(context.Background(), &CreateUserInput{
Email: "bad@test.com",
Password: "strong-pass",
Role: "superuser",
})
require.Error(t, err)
require.Empty(t, repo.created, "非法角色不应写入用户")
}
func TestAdminService_UpdateUser_PromoteToAdmin(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}}
repo := &rpmUserRepoStub{userRepoStub: base}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &redeemRepoStub{},
authCacheInvalidator: invalidator,
}
updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleAdmin})
require.NoError(t, err)
require.Equal(t, RoleAdmin, updated.Role)
require.Equal(t, []int64{42}, invalidator.userIDs, "角色变更应失效认证缓存")
}
func TestAdminService_UpdateUser_RoleOmittedKeepsExisting(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleAdmin}}
repo := &rpmUserRepoStub{userRepoStub: base}
svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}}
newName := "renamed"
updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Username: &newName})
require.NoError(t, err)
require.Equal(t, RoleAdmin, updated.Role, "未提供 role 时不应改变现有角色")
}
func TestAdminService_UpdateUser_InvalidRoleRejected(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}}
repo := &rpmUserRepoStub{userRepoStub: base}
svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}}
_, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: "root"})
require.Error(t, err)
require.Nil(t, repo.lastUpdated, "非法角色不应触发持久化")
}
// roleGuardUserRepoStub 在 rpmUserRepoStub 之上提供可控的管理员计数,
// 用于测试"最后一个管理员不可降级"守卫。
type roleGuardUserRepoStub struct {
*rpmUserRepoStub
adminTotal int64
listCalls int
}
func (s *roleGuardUserRepoStub) ListWithFilters(_ context.Context, _ pagination.PaginationParams, _ UserListFilters) ([]User, *pagination.PaginationResult, error) {
s.listCalls++
return nil, &pagination.PaginationResult{Total: s.adminTotal}, nil
}
func TestAdminService_UpdateUser_DemoteLastAdminRejected(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "a@example.com", Role: RoleAdmin}}
repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 1}
svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}}
_, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleUser})
require.Error(t, err)
require.Contains(t, err.Error(), "last admin")
require.Nil(t, repo.lastUpdated, "最后一个管理员不应被降级持久化")
require.Equal(t, 1, repo.listCalls, "降级路径应触发管理员计数")
}
func TestAdminService_UpdateUser_DemoteAdminAllowedWhenOthersExist(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "a@example.com", Role: RoleAdmin}}
repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 2}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &redeemRepoStub{},
authCacheInvalidator: invalidator,
}
updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleUser})
require.NoError(t, err)
require.Equal(t, RoleUser, updated.Role)
require.NotNil(t, repo.lastUpdated)
require.Equal(t, RoleUser, repo.lastUpdated.Role, "存在其他管理员时允许降级")
}
func TestAdminService_UpdateUser_PromoteDoesNotCountAdmins(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}}
repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 1}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &redeemRepoStub{},
authCacheInvalidator: &authCacheInvalidatorStub{},
}
updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleAdmin})
require.NoError(t, err)
require.Equal(t, RoleAdmin, updated.Role)
require.Equal(t, 0, repo.listCalls, "升级路径不应触发管理员计数")
}
@@ -0,0 +1,112 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type rpmStatusUserRepoStub struct {
UserRepository
user *User
}
func (s *rpmStatusUserRepoStub) GetByID(_ context.Context, _ int64) (*User, error) {
return s.user, nil
}
type rpmStatusAPIKeyRepoStub struct {
APIKeyRepository
keys []APIKey
}
func (s *rpmStatusAPIKeyRepoStub) ListByUserID(_ context.Context, _ int64, _ pagination.PaginationParams, _ APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
return s.keys, &pagination.PaginationResult{Total: int64(len(s.keys))}, nil
}
type rpmStatusGroupRepoStub struct {
GroupRepository
groups map[int64]*Group
}
func (s *rpmStatusGroupRepoStub) GetByIDLite(_ context.Context, id int64) (*Group, error) {
return s.groups[id], nil
}
type rpmStatusRateRepoStub struct {
UserGroupRateRepository
overrides map[int64]*int
}
func (s *rpmStatusRateRepoStub) GetRPMOverrideByUserAndGroup(_ context.Context, _, groupID int64) (*int, error) {
return s.overrides[groupID], nil
}
type rpmStatusCacheStub struct {
UserRPMCache
userUsed int
groupUsed map[int64]int
}
func (s *rpmStatusCacheStub) IncrementUserGroupRPM(context.Context, int64, int64) (int, error) {
return 0, nil
}
func (s *rpmStatusCacheStub) IncrementUserRPM(context.Context, int64) (int, error) {
return 0, nil
}
func (s *rpmStatusCacheStub) GetUserGroupRPM(_ context.Context, _, groupID int64) (int, error) {
return s.groupUsed[groupID], nil
}
func (s *rpmStatusCacheStub) GetUserRPM(context.Context, int64) (int, error) {
return s.userUsed, nil
}
func TestAdminService_GetUserRPMStatus_AggregatesUserAndGroupLimits(t *testing.T) {
groupOneID := int64(1)
groupTwoID := int64(2)
override := 7
svc := &adminServiceImpl{
userRepo: &rpmStatusUserRepoStub{user: &User{
ID: 42,
RPMLimit: 20,
}},
apiKeyRepo: &rpmStatusAPIKeyRepoStub{keys: []APIKey{
{ID: 100, UserID: 42, GroupID: &groupTwoID},
{ID: 101, UserID: 42, GroupID: &groupOneID},
{ID: 102, UserID: 42, GroupID: &groupTwoID},
{ID: 103, UserID: 42},
}},
groupRepo: &rpmStatusGroupRepoStub{groups: map[int64]*Group{
groupOneID: {ID: groupOneID, Name: "group-one", RPMLimit: 10},
groupTwoID: {ID: groupTwoID, Name: "group-two", RPMLimit: 60},
}},
userGroupRateRepo: &rpmStatusRateRepoStub{overrides: map[int64]*int{
groupTwoID: &override,
}},
userRPMCache: &rpmStatusCacheStub{
userUsed: 5,
groupUsed: map[int64]int{
groupOneID: 3,
groupTwoID: 4,
},
},
}
status, err := svc.GetUserRPMStatus(context.Background(), 42)
require.NoError(t, err)
require.Equal(t, &UserRPMStatus{
UserRPMUsed: 5,
UserRPMLimit: 20,
PerGroup: []UserGroupRPMStatus{
{GroupID: groupOneID, GroupName: "group-one", Used: 3, Limit: 10, Source: "group"},
{GroupID: groupTwoID, GroupName: "group-two", Used: 4, Limit: 7, Source: "override"},
},
}, status)
}
@@ -0,0 +1,268 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type accountRepoStubForAdminList struct {
accountRepoStub
listWithFiltersCalls int
listWithFiltersParams pagination.PaginationParams
listWithFiltersPlatform string
listWithFiltersType string
listWithFiltersStatus string
listWithFiltersSearch string
listWithFiltersPrivacy string
listWithFiltersAccounts []Account
listWithFiltersResult *pagination.PaginationResult
listWithFiltersErr error
}
func (s *accountRepoStubForAdminList) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) {
return nil, nil
}
func (s *accountRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
s.listWithFiltersCalls++
s.listWithFiltersParams = params
s.listWithFiltersPlatform = platform
s.listWithFiltersType = accountType
s.listWithFiltersStatus = status
s.listWithFiltersSearch = search
s.listWithFiltersPrivacy = privacyMode
if s.listWithFiltersErr != nil {
return nil, nil, s.listWithFiltersErr
}
result := s.listWithFiltersResult
if result == nil {
result = &pagination.PaginationResult{
Total: int64(len(s.listWithFiltersAccounts)),
Page: params.Page,
PageSize: params.PageSize,
}
}
return s.listWithFiltersAccounts, result, nil
}
type proxyRepoStubForAdminList struct {
proxyRepoStub
listWithFiltersCalls int
listWithFiltersParams pagination.PaginationParams
listWithFiltersProtocol string
listWithFiltersStatus string
listWithFiltersSearch string
listWithFiltersProxies []Proxy
listWithFiltersResult *pagination.PaginationResult
listWithFiltersErr error
listWithFiltersAndAccountCountCalls int
listWithFiltersAndAccountCountParams pagination.PaginationParams
listWithFiltersAndAccountCountProtocol string
listWithFiltersAndAccountCountStatus string
listWithFiltersAndAccountCountSearch string
listWithFiltersAndAccountCountProxies []ProxyWithAccountCount
listWithFiltersAndAccountCountResult *pagination.PaginationResult
listWithFiltersAndAccountCountErr error
}
func (s *proxyRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, protocol, status, search string) ([]Proxy, *pagination.PaginationResult, error) {
s.listWithFiltersCalls++
s.listWithFiltersParams = params
s.listWithFiltersProtocol = protocol
s.listWithFiltersStatus = status
s.listWithFiltersSearch = search
if s.listWithFiltersErr != nil {
return nil, nil, s.listWithFiltersErr
}
result := s.listWithFiltersResult
if result == nil {
result = &pagination.PaginationResult{
Total: int64(len(s.listWithFiltersProxies)),
Page: params.Page,
PageSize: params.PageSize,
}
}
return s.listWithFiltersProxies, result, nil
}
func (s *proxyRepoStubForAdminList) ListWithFiltersAndAccountCount(_ context.Context, params pagination.PaginationParams, protocol, status, search string) ([]ProxyWithAccountCount, *pagination.PaginationResult, error) {
s.listWithFiltersAndAccountCountCalls++
s.listWithFiltersAndAccountCountParams = params
s.listWithFiltersAndAccountCountProtocol = protocol
s.listWithFiltersAndAccountCountStatus = status
s.listWithFiltersAndAccountCountSearch = search
if s.listWithFiltersAndAccountCountErr != nil {
return nil, nil, s.listWithFiltersAndAccountCountErr
}
result := s.listWithFiltersAndAccountCountResult
if result == nil {
result = &pagination.PaginationResult{
Total: int64(len(s.listWithFiltersAndAccountCountProxies)),
Page: params.Page,
PageSize: params.PageSize,
}
}
return s.listWithFiltersAndAccountCountProxies, result, nil
}
type redeemRepoStubForAdminList struct {
redeemRepoStub
listWithFiltersCalls int
listWithFiltersParams pagination.PaginationParams
listWithFiltersType string
listWithFiltersStatus string
listWithFiltersSearch string
listWithFiltersCodes []RedeemCode
listWithFiltersResult *pagination.PaginationResult
listWithFiltersErr error
}
func (s *redeemRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, codeType, status, search string) ([]RedeemCode, *pagination.PaginationResult, error) {
s.listWithFiltersCalls++
s.listWithFiltersParams = params
s.listWithFiltersType = codeType
s.listWithFiltersStatus = status
s.listWithFiltersSearch = search
if s.listWithFiltersErr != nil {
return nil, nil, s.listWithFiltersErr
}
result := s.listWithFiltersResult
if result == nil {
result = &pagination.PaginationResult{
Total: int64(len(s.listWithFiltersCodes)),
Page: params.Page,
PageSize: params.PageSize,
}
}
return s.listWithFiltersCodes, result, nil
}
func (s *redeemRepoStubForAdminList) ListByUserPaginated(_ context.Context, userID int64, params pagination.PaginationParams, codeType string) ([]RedeemCode, *pagination.PaginationResult, error) {
panic("unexpected ListByUserPaginated call")
}
func (s *redeemRepoStubForAdminList) SumPositiveBalanceByUser(_ context.Context, userID int64) (float64, error) {
panic("unexpected SumPositiveBalanceByUser call")
}
func TestAdminService_ListAccounts_WithSearch(t *testing.T) {
t.Run("search 参数正常传递到 repository 层", func(t *testing.T) {
repo := &accountRepoStubForAdminList{
listWithFiltersAccounts: []Account{{ID: 1, Name: "acc"}},
listWithFiltersResult: &pagination.PaginationResult{Total: 10},
}
svc := &adminServiceImpl{accountRepo: repo}
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformGemini, AccountTypeOAuth, StatusActive, "acc", 0, "", "name", "ASC")
require.NoError(t, err)
require.Equal(t, int64(10), total)
require.Equal(t, []Account{{ID: 1, Name: "acc"}}, accounts)
require.Equal(t, 1, repo.listWithFiltersCalls)
require.Equal(t, pagination.PaginationParams{Page: 1, PageSize: 20, SortBy: "name", SortOrder: "ASC"}, repo.listWithFiltersParams)
require.Equal(t, PlatformGemini, repo.listWithFiltersPlatform)
require.Equal(t, AccountTypeOAuth, repo.listWithFiltersType)
require.Equal(t, StatusActive, repo.listWithFiltersStatus)
require.Equal(t, "acc", repo.listWithFiltersSearch)
})
}
func TestAdminService_ListAccounts_WithPrivacyMode(t *testing.T) {
t.Run("privacy_mode 参数正常传递到 repository 层", func(t *testing.T) {
repo := &accountRepoStubForAdminList{
listWithFiltersAccounts: []Account{{ID: 2, Name: "acc2"}},
listWithFiltersResult: &pagination.PaginationResult{Total: 1},
}
svc := &adminServiceImpl{accountRepo: repo}
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformOpenAI, AccountTypeOAuth, StatusActive, "acc2", 0, PrivacyModeCFBlocked, "", "")
require.NoError(t, err)
require.Equal(t, int64(1), total)
require.Equal(t, []Account{{ID: 2, Name: "acc2"}}, accounts)
require.Equal(t, PrivacyModeCFBlocked, repo.listWithFiltersPrivacy)
})
}
func TestAdminService_ListProxies_WithSearch(t *testing.T) {
t.Run("search 参数正常传递到 repository 层", func(t *testing.T) {
repo := &proxyRepoStubForAdminList{
listWithFiltersProxies: []Proxy{{ID: 2, Name: "p1"}},
listWithFiltersResult: &pagination.PaginationResult{Total: 7},
}
svc := &adminServiceImpl{proxyRepo: repo}
proxies, total, err := svc.ListProxies(context.Background(), 3, 50, "http", StatusActive, "p1", "name", "ASC")
require.NoError(t, err)
require.Equal(t, int64(7), total)
require.Equal(t, []Proxy{{ID: 2, Name: "p1"}}, proxies)
require.Equal(t, 1, repo.listWithFiltersCalls)
require.Equal(t, pagination.PaginationParams{Page: 3, PageSize: 50, SortBy: "name", SortOrder: "ASC"}, repo.listWithFiltersParams)
require.Equal(t, "http", repo.listWithFiltersProtocol)
require.Equal(t, StatusActive, repo.listWithFiltersStatus)
require.Equal(t, "p1", repo.listWithFiltersSearch)
})
}
func TestAdminService_ListProxiesWithAccountCount_WithSearch(t *testing.T) {
t.Run("search 参数正常传递到 repository 层", func(t *testing.T) {
repo := &proxyRepoStubForAdminList{
listWithFiltersAndAccountCountProxies: []ProxyWithAccountCount{{Proxy: Proxy{ID: 3, Name: "p2"}, AccountCount: 5}},
listWithFiltersAndAccountCountResult: &pagination.PaginationResult{Total: 9},
}
svc := &adminServiceImpl{proxyRepo: repo}
proxies, total, err := svc.ListProxiesWithAccountCount(context.Background(), 2, 10, "socks5", StatusDisabled, "p2", "account_count", "DESC")
require.NoError(t, err)
require.Equal(t, int64(9), total)
require.Equal(t, []ProxyWithAccountCount{{Proxy: Proxy{ID: 3, Name: "p2"}, AccountCount: 5}}, proxies)
require.Equal(t, 1, repo.listWithFiltersAndAccountCountCalls)
require.Equal(t, pagination.PaginationParams{Page: 2, PageSize: 10, SortBy: "account_count", SortOrder: "DESC"}, repo.listWithFiltersAndAccountCountParams)
require.Equal(t, "socks5", repo.listWithFiltersAndAccountCountProtocol)
require.Equal(t, StatusDisabled, repo.listWithFiltersAndAccountCountStatus)
require.Equal(t, "p2", repo.listWithFiltersAndAccountCountSearch)
})
}
func TestAdminService_ListRedeemCodes_WithSearch(t *testing.T) {
t.Run("search 参数正常传递到 repository 层", func(t *testing.T) {
repo := &redeemRepoStubForAdminList{
listWithFiltersCodes: []RedeemCode{{ID: 4, Code: "ABC"}},
listWithFiltersResult: &pagination.PaginationResult{Total: 3},
}
svc := &adminServiceImpl{redeemCodeRepo: repo}
codes, total, err := svc.ListRedeemCodes(context.Background(), 1, 20, RedeemTypeBalance, StatusUnused, "ABC", "value", "ASC")
require.NoError(t, err)
require.Equal(t, int64(3), total)
require.Equal(t, []RedeemCode{{ID: 4, Code: "ABC"}}, codes)
require.Equal(t, 1, repo.listWithFiltersCalls)
require.Equal(t, pagination.PaginationParams{Page: 1, PageSize: 20, SortBy: "value", SortOrder: "ASC"}, repo.listWithFiltersParams)
require.Equal(t, RedeemTypeBalance, repo.listWithFiltersType)
require.Equal(t, StatusUnused, repo.listWithFiltersStatus)
require.Equal(t, "ABC", repo.listWithFiltersSearch)
})
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,263 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/require"
)
type balanceUserRepoStub struct {
*userRepoStub
adjustErr error
// changes 记录每次原子余额变更,顺序与调用顺序一致。
changes []BalanceChange
}
func (s *balanceUserRepoStub) AdjustBalance(ctx context.Context, id int64, delta float64) (BalanceChange, error) {
return s.apply(func(current float64) float64 { return current + delta })
}
func (s *balanceUserRepoStub) SetBalance(ctx context.Context, id int64, value float64) (BalanceChange, error) {
return s.apply(func(float64) float64 { return value })
}
func (s *balanceUserRepoStub) apply(next func(current float64) float64) (BalanceChange, error) {
if s.adjustErr != nil {
return BalanceChange{}, s.adjustErr
}
if s.userRepoStub == nil || s.userRepoStub.user == nil {
return BalanceChange{}, ErrUserNotFound
}
change := BalanceChange{Old: s.userRepoStub.user.Balance}
change.New = next(change.Old)
if change.New < 0 {
return change, ErrBalanceNegative
}
s.userRepoStub.user.Balance = change.New
s.changes = append(s.changes, change)
return change, nil
}
type balanceRedeemRepoStub struct {
*redeemRepoStub
created []*RedeemCode
}
func (s *balanceRedeemRepoStub) Create(ctx context.Context, code *RedeemCode) error {
if code == nil {
return nil
}
clone := *code
s.created = append(s.created, &clone)
return nil
}
type authCacheInvalidatorStub struct {
userIDs []int64
groupIDs []int64
keys []string
}
type adminRechargeAffiliateAccruerStub struct {
calls []adminRechargeAffiliateAccrual
rebate float64
err error
}
type adminRechargeAffiliateAccrual struct {
userID int64
amount float64
}
func (s *adminRechargeAffiliateAccruerStub) AccrueInviteRebate(_ context.Context, userID int64, amount float64) (float64, error) {
s.calls = append(s.calls, adminRechargeAffiliateAccrual{userID: userID, amount: amount})
return s.rebate, s.err
}
func adminRechargeSettingService(enabled bool) *SettingService {
values := map[string]string{}
if enabled {
values[SettingKeyAffiliateAdminRechargeEnabled] = "true"
}
return NewSettingService(&settingRepoStub{values: values}, nil)
}
func (s *authCacheInvalidatorStub) InvalidateAuthCacheByKey(ctx context.Context, key string) {
s.keys = append(s.keys, key)
}
func (s *authCacheInvalidatorStub) InvalidateAuthCacheByUserID(ctx context.Context, userID int64) {
s.userIDs = append(s.userIDs, userID)
}
func (s *authCacheInvalidatorStub) InvalidateAuthCacheByGroupID(ctx context.Context, groupID int64) {
s.groupIDs = append(s.groupIDs, groupID)
}
// 管理员调账必须走原子的 AdjustBalance/SetBalance,而不是"读余额→算新值→整行写回",
// 后者会把并发的计费扣款覆盖掉。userRepoStub.Update 对未预期的调用会 panic
// 因此这里同时证明它没被走到。
func TestAdminService_UpdateUserBalance_UsesAtomicPrimitives(t *testing.T) {
tests := []struct {
name string
operation string
amount float64
want BalanceChange
}{
{name: "add", operation: "add", amount: 5, want: BalanceChange{Old: 10, New: 15}},
{name: "subtract", operation: "subtract", amount: 4, want: BalanceChange{Old: 10, New: 6}},
{name: "set", operation: "set", amount: 2, want: BalanceChange{Old: 10, New: 2}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &balanceUserRepoStub{userRepoStub: &userRepoStub{user: &User{ID: 7, Balance: 10}}}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}},
}
user, err := svc.UpdateUserBalance(context.Background(), 7, tt.amount, tt.operation, "")
require.NoError(t, err)
require.Equal(t, []BalanceChange{tt.want}, repo.changes)
require.Equal(t, tt.want.New, user.Balance)
})
}
}
func TestAdminService_UpdateUserBalance_RejectsNegativeResult(t *testing.T) {
repo := &balanceUserRepoStub{userRepoStub: &userRepoStub{user: &User{ID: 7, Balance: 3}}}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}},
}
_, err := svc.UpdateUserBalance(context.Background(), 7, 4, "subtract", "")
require.Error(t, err)
require.Contains(t, err.Error(), "balance cannot be negative")
require.Empty(t, repo.changes, "refused adjustment must not be applied")
require.Equal(t, 3.0, repo.userRepoStub.user.Balance)
}
func TestAdminService_UpdateUserBalance_RejectsUnknownOperation(t *testing.T) {
repo := &balanceUserRepoStub{userRepoStub: &userRepoStub{user: &User{ID: 7, Balance: 10}}}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}},
}
_, err := svc.UpdateUserBalance(context.Background(), 7, 1, "multiply", "")
require.Error(t, err)
require.Empty(t, repo.changes)
}
func TestAdminService_UpdateUserBalance_InvalidatesAuthCache(t *testing.T) {
baseRepo := &userRepoStub{user: &User{ID: 7, Balance: 10}}
repo := &balanceUserRepoStub{userRepoStub: baseRepo}
redeemRepo := &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: redeemRepo,
authCacheInvalidator: invalidator,
}
_, err := svc.UpdateUserBalance(context.Background(), 7, 5, "add", "")
require.NoError(t, err)
require.Equal(t, []int64{7}, invalidator.userIDs)
require.Len(t, redeemRepo.created, 1)
}
func TestAdminService_UpdateUserBalance_NoChangeNoInvalidate(t *testing.T) {
baseRepo := &userRepoStub{user: &User{ID: 7, Balance: 10}}
repo := &balanceUserRepoStub{userRepoStub: baseRepo}
redeemRepo := &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: redeemRepo,
authCacheInvalidator: invalidator,
}
_, err := svc.UpdateUserBalance(context.Background(), 7, 10, "set", "")
require.NoError(t, err)
require.Empty(t, invalidator.userIDs)
require.Empty(t, redeemRepo.created)
}
func TestAdminService_UpdateUserBalance_AdminRechargeAffiliateRebate(t *testing.T) {
tests := []struct {
name string
enabled bool
operation string
amount float64
wantCalls []adminRechargeAffiliateAccrual
}{
{
name: "disabled by default",
operation: "add",
amount: 5,
},
{
name: "enabled add",
enabled: true,
operation: "add",
amount: 0.1,
wantCalls: []adminRechargeAffiliateAccrual{{userID: 7, amount: 0.1}},
},
{
name: "enabled set increase",
enabled: true,
operation: "set",
amount: 15,
},
{
name: "enabled subtract",
enabled: true,
operation: "subtract",
amount: 5,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
baseRepo := &userRepoStub{user: &User{ID: 7, Balance: 10}}
repo := &balanceUserRepoStub{userRepoStub: baseRepo}
redeemRepo := &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}}
affiliate := &adminRechargeAffiliateAccruerStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: redeemRepo,
settingService: adminRechargeSettingService(tt.enabled),
affiliateService: affiliate,
}
_, err := svc.UpdateUserBalance(context.Background(), 7, tt.amount, tt.operation, "")
require.NoError(t, err)
require.Equal(t, tt.wantCalls, affiliate.calls)
})
}
}
func TestAdminService_UpdateUserBalance_AffiliateFailureDoesNotRollbackRecharge(t *testing.T) {
baseRepo := &userRepoStub{user: &User{ID: 7, Balance: 10}}
repo := &balanceUserRepoStub{userRepoStub: baseRepo}
redeemRepo := &balanceRedeemRepoStub{redeemRepoStub: &redeemRepoStub{}}
affiliate := &adminRechargeAffiliateAccruerStub{err: errors.New("affiliate unavailable")}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: redeemRepo,
settingService: adminRechargeSettingService(true),
affiliateService: affiliate,
}
user, err := svc.UpdateUserBalance(context.Background(), 7, 5, "add", "")
require.NoError(t, err)
require.Equal(t, 15.0, user.Balance)
require.Equal(t, []adminRechargeAffiliateAccrual{{userID: 7, amount: 5}}, affiliate.calls)
require.Len(t, redeemRepo.created, 1)
}
@@ -0,0 +1,69 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
// rpmUserRepoStub 复用 admin_service_update_balance_test.go 的基础 stub 结构,
// 只在 Update 时把入参克隆一份,便于断言修改后的 RPMLimit。
type rpmUserRepoStub struct {
*userRepoStub
lastUpdated *User
}
func (s *rpmUserRepoStub) Update(_ context.Context, user *User, _ UserUpdateFields) error {
if user == nil {
return nil
}
clone := *user
s.lastUpdated = &clone
if s.userRepoStub != nil {
s.userRepoStub.user = &clone
}
return nil
}
func TestAdminService_UpdateUser_InvalidatesAuthCacheOnRPMLimitChange(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", RPMLimit: 10}}
repo := &rpmUserRepoStub{userRepoStub: base}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &redeemRepoStub{},
authCacheInvalidator: invalidator,
}
newRPM := 60
updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{
RPMLimit: &newRPM,
})
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 60, updated.RPMLimit)
require.Equal(t, []int64{42}, invalidator.userIDs, "仅修改 RPMLimit 也应失效 API Key 认证缓存")
}
func TestAdminService_UpdateUser_NoInvalidateWhenRPMLimitUnchanged(t *testing.T) {
base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", RPMLimit: 10, Username: "old"}}
repo := &rpmUserRepoStub{userRepoStub: base}
invalidator := &authCacheInvalidatorStub{}
svc := &adminServiceImpl{
userRepo: repo,
redeemCodeRepo: &redeemRepoStub{},
authCacheInvalidator: invalidator,
}
newName := "new"
sameRPM := 10
_, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{
Username: &newName,
RPMLimit: &sameRPM,
})
require.NoError(t, err)
require.Empty(t, invalidator.userIDs, "只改 username 不应触发认证缓存失效")
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,624 @@
package service
import (
"context"
"errors"
"math"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)
var (
ErrAffiliateProfileNotFound = infraerrors.NotFound("AFFILIATE_PROFILE_NOT_FOUND", "affiliate profile not found")
ErrAffiliateCodeInvalid = infraerrors.BadRequest("AFFILIATE_CODE_INVALID", "invalid affiliate code")
ErrAffiliateCodeTaken = infraerrors.Conflict("AFFILIATE_CODE_TAKEN", "affiliate code already in use")
ErrAffiliateAlreadyBound = infraerrors.Conflict("AFFILIATE_ALREADY_BOUND", "affiliate inviter already bound")
ErrAffiliateQuotaEmpty = infraerrors.BadRequest("AFFILIATE_QUOTA_EMPTY", "no affiliate quota available to transfer")
)
const (
affiliateInviteesLimit = 100
// AffiliateCodeMinLength / AffiliateCodeMaxLength bound both system-generated
// 12-char codes and admin-customized codes (e.g. "VIP2026").
AffiliateCodeMinLength = 4
AffiliateCodeMaxLength = 32
)
// affiliateCodeValidChar accepts uppercase letters, digits, underscore and dash.
// All input passes through strings.ToUpper before validation, so lowercase from
// users is normalized — admins may supply mixed case in their UI.
var affiliateCodeValidChar = func() [256]bool {
var tbl [256]bool
for c := byte('A'); c <= 'Z'; c++ {
tbl[c] = true
}
for c := byte('0'); c <= '9'; c++ {
tbl[c] = true
}
tbl['_'] = true
tbl['-'] = true
return tbl
}()
// isValidAffiliateCodeFormat validates code format for both binding (user input)
// and admin updates. Caller is expected to upper-case the input first.
func isValidAffiliateCodeFormat(code string) bool {
if len(code) < AffiliateCodeMinLength || len(code) > AffiliateCodeMaxLength {
return false
}
for i := 0; i < len(code); i++ {
if !affiliateCodeValidChar[code[i]] {
return false
}
}
return true
}
type AffiliateSummary struct {
UserID int64 `json:"user_id"`
AffCode string `json:"aff_code"`
AffCodeCustom bool `json:"aff_code_custom"`
AffRebateRatePercent *float64 `json:"aff_rebate_rate_percent,omitempty"`
InviterID *int64 `json:"inviter_id,omitempty"`
AffCount int `json:"aff_count"`
AffQuota float64 `json:"aff_quota"`
AffFrozenQuota float64 `json:"aff_frozen_quota"`
AffHistoryQuota float64 `json:"aff_history_quota"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type AffiliateInvitee struct {
UserID int64 `json:"user_id"`
Email string `json:"email"`
Username string `json:"username"`
CreatedAt *time.Time `json:"created_at,omitempty"`
TotalRebate float64 `json:"total_rebate"`
}
type AffiliateDetail struct {
UserID int64 `json:"user_id"`
AffCode string `json:"aff_code"`
InviterID *int64 `json:"inviter_id,omitempty"`
AffCount int `json:"aff_count"`
AffQuota float64 `json:"aff_quota"`
AffFrozenQuota float64 `json:"aff_frozen_quota"`
AffHistoryQuota float64 `json:"aff_history_quota"`
// EffectiveRebateRatePercent 是当前用户作为邀请人时实际生效的返利比例:
// 优先用户自己的专属比例(aff_rebate_rate_percent),否则回退到全局比例。
// 用于在用户的 /affiliate 页面直观展示「分享后能拿到多少」。
EffectiveRebateRatePercent float64 `json:"effective_rebate_rate_percent"`
Invitees []AffiliateInvitee `json:"invitees"`
}
type AffiliateRepository interface {
EnsureUserAffiliate(ctx context.Context, userID int64) (*AffiliateSummary, error)
GetAffiliateByCode(ctx context.Context, code string) (*AffiliateSummary, error)
BindInviter(ctx context.Context, userID, inviterID int64) (bool, error)
AccrueQuota(ctx context.Context, inviterID, inviteeUserID int64, amount float64, freezeHours int, sourceOrderID *int64) (bool, error)
GetAccruedRebateFromInvitee(ctx context.Context, inviterID, inviteeUserID int64) (float64, error)
ThawFrozenQuota(ctx context.Context, userID int64) (float64, error)
TransferQuotaToBalance(ctx context.Context, userID int64) (float64, float64, error)
ListInvitees(ctx context.Context, inviterID int64, limit int) ([]AffiliateInvitee, error)
// 管理端:用户级专属配置
UpdateUserAffCode(ctx context.Context, userID int64, newCode string) error
ResetUserAffCode(ctx context.Context, userID int64) (string, error)
SetUserRebateRate(ctx context.Context, userID int64, ratePercent *float64) error
BatchSetUserRebateRate(ctx context.Context, userIDs []int64, ratePercent *float64) error
ListUsersWithCustomSettings(ctx context.Context, filter AffiliateAdminFilter) ([]AffiliateAdminEntry, int64, error)
ListAffiliateInviteRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateInviteRecord, int64, error)
ListAffiliateRebateRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateRebateRecord, int64, error)
ListAffiliateTransferRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateTransferRecord, int64, error)
GetAffiliateUserOverview(ctx context.Context, userID int64) (*AffiliateUserOverview, error)
}
// AffiliateAdminFilter 列表筛选条件
type AffiliateAdminFilter struct {
Search string
Page int
PageSize int
}
// AffiliateAdminEntry 专属用户列表条目
type AffiliateAdminEntry struct {
UserID int64 `json:"user_id"`
Email string `json:"email"`
Username string `json:"username"`
AffCode string `json:"aff_code"`
AffCodeCustom bool `json:"aff_code_custom"`
AffRebateRatePercent *float64 `json:"aff_rebate_rate_percent,omitempty"`
AffCount int `json:"aff_count"`
}
type AffiliateRecordFilter struct {
Search string
Page int
PageSize int
StartAt *time.Time
EndAt *time.Time
SortBy string
SortDesc bool
}
type AffiliateInviteRecord struct {
InviterID int64 `json:"inviter_id"`
InviterEmail string `json:"inviter_email"`
InviterUsername string `json:"inviter_username"`
InviteeID int64 `json:"invitee_id"`
InviteeEmail string `json:"invitee_email"`
InviteeUsername string `json:"invitee_username"`
AffCode string `json:"aff_code"`
TotalRebate float64 `json:"total_rebate"`
CreatedAt time.Time `json:"created_at"`
}
type AffiliateRebateRecord struct {
OrderID int64 `json:"order_id"`
OutTradeNo string `json:"out_trade_no"`
InviterID int64 `json:"inviter_id"`
InviterEmail string `json:"inviter_email"`
InviterUsername string `json:"inviter_username"`
InviteeID int64 `json:"invitee_id"`
InviteeEmail string `json:"invitee_email"`
InviteeUsername string `json:"invitee_username"`
OrderAmount float64 `json:"order_amount"`
PayAmount float64 `json:"pay_amount"`
RebateAmount float64 `json:"rebate_amount"`
PaymentType string `json:"payment_type"`
OrderStatus string `json:"order_status"`
CreatedAt time.Time `json:"created_at"`
}
type AffiliateTransferRecord struct {
LedgerID int64 `json:"ledger_id"`
UserID int64 `json:"user_id"`
UserEmail string `json:"user_email"`
Username string `json:"username"`
Amount float64 `json:"amount"`
BalanceAfter *float64 `json:"balance_after,omitempty"`
AvailableQuotaAfter *float64 `json:"available_quota_after,omitempty"`
FrozenQuotaAfter *float64 `json:"frozen_quota_after,omitempty"`
HistoryQuotaAfter *float64 `json:"history_quota_after,omitempty"`
SnapshotAvailable bool `json:"snapshot_available"`
CurrentBalance float64 `json:"-"`
RemainingQuota float64 `json:"-"`
FrozenQuota float64 `json:"-"`
HistoryQuota float64 `json:"-"`
CreatedAt time.Time `json:"created_at"`
}
type AffiliateUserOverview struct {
UserID int64 `json:"user_id"`
Email string `json:"email"`
Username string `json:"username"`
AffCode string `json:"aff_code"`
RebateRatePercent float64 `json:"rebate_rate_percent"`
RebateRateCustom bool `json:"-"`
InvitedCount int `json:"invited_count"`
RebatedInviteeCount int `json:"rebated_invitee_count"`
AvailableQuota float64 `json:"available_quota"`
HistoryQuota float64 `json:"history_quota"`
}
type AffiliateService struct {
repo AffiliateRepository
settingService *SettingService
authCacheInvalidator APIKeyAuthCacheInvalidator
billingCacheService *BillingCacheService
}
func NewAffiliateService(repo AffiliateRepository, settingService *SettingService, authCacheInvalidator APIKeyAuthCacheInvalidator, billingCacheService *BillingCacheService) *AffiliateService {
return &AffiliateService{
repo: repo,
settingService: settingService,
authCacheInvalidator: authCacheInvalidator,
billingCacheService: billingCacheService,
}
}
// IsEnabled reports whether the affiliate (邀请返利) feature is turned on.
func (s *AffiliateService) IsEnabled(ctx context.Context) bool {
if s == nil || s.settingService == nil {
return AffiliateEnabledDefault
}
return s.settingService.IsAffiliateEnabled(ctx)
}
func (s *AffiliateService) EnsureUserAffiliate(ctx context.Context, userID int64) (*AffiliateSummary, error) {
if userID <= 0 {
return nil, infraerrors.BadRequest("INVALID_USER", "invalid user")
}
if s == nil || s.repo == nil {
return nil, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.EnsureUserAffiliate(ctx, userID)
}
func (s *AffiliateService) GetAffiliateDetail(ctx context.Context, userID int64) (*AffiliateDetail, error) {
// Lazy thaw: move any matured frozen quota to available before reading.
if s != nil && s.repo != nil {
// best-effort: thaw failure is non-fatal
_, _ = s.repo.ThawFrozenQuota(ctx, userID)
}
summary, err := s.EnsureUserAffiliate(ctx, userID)
if err != nil {
return nil, err
}
invitees, err := s.listInvitees(ctx, userID)
if err != nil {
return nil, err
}
return &AffiliateDetail{
UserID: summary.UserID,
AffCode: summary.AffCode,
InviterID: summary.InviterID,
AffCount: summary.AffCount,
AffQuota: summary.AffQuota,
AffFrozenQuota: summary.AffFrozenQuota,
AffHistoryQuota: summary.AffHistoryQuota,
EffectiveRebateRatePercent: s.resolveRebateRatePercent(ctx, summary),
Invitees: invitees,
}, nil
}
func (s *AffiliateService) BindInviterByCode(ctx context.Context, userID int64, rawCode string) error {
code := strings.ToUpper(strings.TrimSpace(rawCode))
if code == "" {
return nil
}
if s == nil || s.repo == nil {
return infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
// 总开关关闭时,注册阶段静默忽略 aff 参数(不报错,避免阻断注册流程)
if !s.IsEnabled(ctx) {
return nil
}
if !isValidAffiliateCodeFormat(code) {
return ErrAffiliateCodeInvalid
}
selfSummary, err := s.repo.EnsureUserAffiliate(ctx, userID)
if err != nil {
return err
}
if selfSummary.InviterID != nil {
return nil
}
inviterSummary, err := s.repo.GetAffiliateByCode(ctx, code)
if err != nil {
if errors.Is(err, ErrAffiliateProfileNotFound) {
return ErrAffiliateCodeInvalid
}
return err
}
if inviterSummary == nil || inviterSummary.UserID <= 0 || inviterSummary.UserID == userID {
return ErrAffiliateCodeInvalid
}
bound, err := s.repo.BindInviter(ctx, userID, inviterSummary.UserID)
if err != nil {
return err
}
if !bound {
return ErrAffiliateAlreadyBound
}
return nil
}
func (s *AffiliateService) AccrueInviteRebate(ctx context.Context, inviteeUserID int64, baseRechargeAmount float64) (float64, error) {
return s.AccrueInviteRebateForOrder(ctx, inviteeUserID, baseRechargeAmount, nil)
}
func (s *AffiliateService) AccrueInviteRebateForOrder(ctx context.Context, inviteeUserID int64, baseRechargeAmount float64, sourceOrderID *int64) (float64, error) {
if s == nil || s.repo == nil {
return 0, nil
}
if inviteeUserID <= 0 || baseRechargeAmount <= 0 || math.IsNaN(baseRechargeAmount) || math.IsInf(baseRechargeAmount, 0) {
return 0, nil
}
// 总开关关闭时,新充值不再产生返利
if !s.IsEnabled(ctx) {
return 0, nil
}
inviteeSummary, err := s.repo.EnsureUserAffiliate(ctx, inviteeUserID)
if err != nil {
return 0, err
}
if inviteeSummary.InviterID == nil || *inviteeSummary.InviterID <= 0 {
return 0, nil
}
// 加载邀请人 profile,优先使用专属比例(覆盖全局)
inviterSummary, err := s.repo.EnsureUserAffiliate(ctx, *inviteeSummary.InviterID)
if err != nil {
return 0, err
}
// 有效期检查:超过返利有效期后不再产生返利
if s.settingService != nil {
if durationDays := s.settingService.GetAffiliateRebateDurationDays(ctx); durationDays > 0 {
if time.Now().After(inviteeSummary.CreatedAt.AddDate(0, 0, durationDays)) {
return 0, nil
}
}
}
rebateRatePercent := s.resolveRebateRatePercent(ctx, inviterSummary)
rebate := roundTo(baseRechargeAmount*(rebateRatePercent/100), 8)
if rebate <= 0 {
return 0, nil
}
// 单人上限检查:精确截断到剩余额度
if s.settingService != nil {
if perInviteeCap := s.settingService.GetAffiliateRebatePerInviteeCap(ctx); perInviteeCap > 0 {
existing, err := s.repo.GetAccruedRebateFromInvitee(ctx, *inviteeSummary.InviterID, inviteeUserID)
if err != nil {
return 0, err
}
if existing >= perInviteeCap {
return 0, nil
}
if remaining := perInviteeCap - existing; rebate > remaining {
rebate = roundTo(remaining, 8)
}
}
}
var freezeHours int
if s.settingService != nil {
freezeHours = s.settingService.GetAffiliateRebateFreezeHours(ctx)
}
applied, err := s.repo.AccrueQuota(ctx, *inviteeSummary.InviterID, inviteeUserID, rebate, freezeHours, sourceOrderID)
if err != nil {
return 0, err
}
if !applied {
return 0, nil
}
return rebate, nil
}
// resolveRebateRatePercent returns the inviter's exclusive rate when set,
// otherwise the global setting value (clamped to [Min, Max]).
func (s *AffiliateService) resolveRebateRatePercent(ctx context.Context, inviter *AffiliateSummary) float64 {
if inviter != nil && inviter.AffRebateRatePercent != nil {
v := *inviter.AffRebateRatePercent
if math.IsNaN(v) || math.IsInf(v, 0) {
return s.globalRebateRatePercent(ctx)
}
return clampAffiliateRebateRate(v)
}
return s.globalRebateRatePercent(ctx)
}
// globalRebateRatePercent reads the system-wide rebate rate via SettingService,
// returning the documented default when SettingService is unavailable.
func (s *AffiliateService) globalRebateRatePercent(ctx context.Context) float64 {
if s == nil || s.settingService == nil {
return AffiliateRebateRateDefault
}
return s.settingService.GetAffiliateRebateRatePercent(ctx)
}
func (s *AffiliateService) TransferAffiliateQuota(ctx context.Context, userID int64) (float64, float64, error) {
if s == nil || s.repo == nil {
return 0, 0, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
transferred, balance, err := s.repo.TransferQuotaToBalance(ctx, userID)
if err != nil {
return 0, 0, err
}
if transferred > 0 {
s.invalidateAffiliateCaches(ctx, userID)
}
return transferred, balance, nil
}
func (s *AffiliateService) listInvitees(ctx context.Context, inviterID int64) ([]AffiliateInvitee, error) {
if s == nil || s.repo == nil {
return nil, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
invitees, err := s.repo.ListInvitees(ctx, inviterID, affiliateInviteesLimit)
if err != nil {
return nil, err
}
for i := range invitees {
invitees[i].Email = maskEmail(invitees[i].Email)
}
return invitees, nil
}
func roundTo(v float64, scale int) float64 {
factor := math.Pow10(scale)
return math.Round(v*factor) / factor
}
func maskEmail(email string) string {
email = strings.TrimSpace(email)
if email == "" {
return ""
}
at := strings.Index(email, "@")
if at <= 0 || at >= len(email)-1 {
return "***"
}
local := email[:at]
domain := email[at+1:]
dot := strings.LastIndex(domain, ".")
maskedLocal := maskSegment(local)
if dot <= 0 || dot >= len(domain)-1 {
return maskedLocal + "@" + maskSegment(domain)
}
domainName := domain[:dot]
tld := domain[dot:]
return maskedLocal + "@" + maskSegment(domainName) + tld
}
func maskSegment(s string) string {
r := []rune(s)
if len(r) == 0 {
return "***"
}
if len(r) == 1 {
return string(r[0]) + "***"
}
return string(r[0]) + "***"
}
func (s *AffiliateService) invalidateAffiliateCaches(ctx context.Context, userID int64) {
if s.authCacheInvalidator != nil {
s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID)
}
if s.billingCacheService != nil {
if err := s.billingCacheService.InvalidateUserBalance(ctx, userID); err != nil {
logger.LegacyPrintf("service.affiliate", "[Affiliate] Failed to invalidate billing cache for user %d: %v", userID, err)
}
}
}
// =========================
// Admin: 专属配置管理
// =========================
// validateExclusiveRate ensures a per-user override is finite and within
// [Min, Max]. nil is always valid (means "clear / fall back to global").
func validateExclusiveRate(ratePercent *float64) error {
if ratePercent == nil {
return nil
}
v := *ratePercent
if math.IsNaN(v) || math.IsInf(v, 0) {
return infraerrors.BadRequest("INVALID_RATE", "invalid rebate rate")
}
if v < AffiliateRebateRateMin || v > AffiliateRebateRateMax {
return infraerrors.BadRequest("INVALID_RATE", "rebate rate out of range")
}
return nil
}
// AdminUpdateUserAffCode 管理员改写用户的邀请码(专属邀请码)。
func (s *AffiliateService) AdminUpdateUserAffCode(ctx context.Context, userID int64, rawCode string) error {
if s == nil || s.repo == nil {
return infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
code := strings.ToUpper(strings.TrimSpace(rawCode))
if !isValidAffiliateCodeFormat(code) {
return ErrAffiliateCodeInvalid
}
return s.repo.UpdateUserAffCode(ctx, userID, code)
}
// AdminResetUserAffCode 重置用户邀请码为系统随机码。
func (s *AffiliateService) AdminResetUserAffCode(ctx context.Context, userID int64) (string, error) {
if s == nil || s.repo == nil {
return "", infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.ResetUserAffCode(ctx, userID)
}
// AdminSetUserRebateRate 设置/清除用户专属返利比例。ratePercent==nil 表示清除。
func (s *AffiliateService) AdminSetUserRebateRate(ctx context.Context, userID int64, ratePercent *float64) error {
if s == nil || s.repo == nil {
return infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
if err := validateExclusiveRate(ratePercent); err != nil {
return err
}
return s.repo.SetUserRebateRate(ctx, userID, ratePercent)
}
// AdminBatchSetUserRebateRate 批量设置/清除用户专属返利比例。
func (s *AffiliateService) AdminBatchSetUserRebateRate(ctx context.Context, userIDs []int64, ratePercent *float64) error {
if s == nil || s.repo == nil {
return infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
if err := validateExclusiveRate(ratePercent); err != nil {
return err
}
cleaned := make([]int64, 0, len(userIDs))
for _, uid := range userIDs {
if uid > 0 {
cleaned = append(cleaned, uid)
}
}
if len(cleaned) == 0 {
return nil
}
return s.repo.BatchSetUserRebateRate(ctx, cleaned, ratePercent)
}
// AdminListCustomUsers 列出有专属配置的用户。
func (s *AffiliateService) AdminListCustomUsers(ctx context.Context, filter AffiliateAdminFilter) ([]AffiliateAdminEntry, int64, error) {
if s == nil || s.repo == nil {
return nil, 0, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.ListUsersWithCustomSettings(ctx, filter)
}
func (s *AffiliateService) AdminListInviteRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateInviteRecord, int64, error) {
if s == nil || s.repo == nil {
return nil, 0, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.ListAffiliateInviteRecords(ctx, normalizeAffiliateRecordFilter(filter))
}
func (s *AffiliateService) AdminListRebateRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateRebateRecord, int64, error) {
if s == nil || s.repo == nil {
return nil, 0, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.ListAffiliateRebateRecords(ctx, normalizeAffiliateRecordFilter(filter))
}
func (s *AffiliateService) AdminListTransferRecords(ctx context.Context, filter AffiliateRecordFilter) ([]AffiliateTransferRecord, int64, error) {
if s == nil || s.repo == nil {
return nil, 0, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
return s.repo.ListAffiliateTransferRecords(ctx, normalizeAffiliateRecordFilter(filter))
}
func (s *AffiliateService) AdminGetUserOverview(ctx context.Context, userID int64) (*AffiliateUserOverview, error) {
if userID <= 0 {
return nil, infraerrors.BadRequest("INVALID_USER", "invalid user")
}
if s == nil || s.repo == nil {
return nil, infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "affiliate service unavailable")
}
overview, err := s.repo.GetAffiliateUserOverview(ctx, userID)
if err != nil {
return nil, err
}
if overview != nil {
if !overview.RebateRateCustom {
overview.RebateRatePercent = s.globalRebateRatePercent(ctx)
}
overview.RebateRatePercent = clampAffiliateRebateRate(overview.RebateRatePercent)
}
return overview, nil
}
func normalizeAffiliateRecordFilter(filter AffiliateRecordFilter) AffiliateRecordFilter {
if filter.Page <= 0 {
filter.Page = 1
}
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.PageSize > 100 {
filter.PageSize = 100
}
filter.Search = strings.TrimSpace(filter.Search)
filter.SortBy = strings.TrimSpace(filter.SortBy)
return filter
}
@@ -0,0 +1,131 @@
//go:build unit
package service
import (
"context"
"math"
"testing"
"github.com/stretchr/testify/require"
)
// TestResolveRebateRatePercent_PerUserOverride verifies that per-inviter
// AffRebateRatePercent overrides the global rate, that NULL falls back to the
// global rate, and that out-of-range exclusive rates are clamped silently.
//
// SettingService is left nil here so globalRebateRatePercent returns the
// documented default (AffiliateRebateRateDefault = 20%) — this exercises the
// fallback path without spinning up a settings stub.
func TestResolveRebateRatePercent_PerUserOverride(t *testing.T) {
t.Parallel()
svc := &AffiliateService{}
// nil exclusive rate → falls back to global default (20%)
require.InDelta(t, AffiliateRebateRateDefault,
svc.resolveRebateRatePercent(context.Background(), &AffiliateSummary{}), 1e-9)
// exclusive rate set → overrides global
rate := 50.0
require.InDelta(t, 50.0,
svc.resolveRebateRatePercent(context.Background(), &AffiliateSummary{AffRebateRatePercent: &rate}), 1e-9)
// exclusive rate 0 → returns 0 (no rebate, intentional)
zero := 0.0
require.InDelta(t, 0.0,
svc.resolveRebateRatePercent(context.Background(), &AffiliateSummary{AffRebateRatePercent: &zero}), 1e-9)
// exclusive rate above max → clamped to Max
tooHigh := 250.0
require.InDelta(t, AffiliateRebateRateMax,
svc.resolveRebateRatePercent(context.Background(), &AffiliateSummary{AffRebateRatePercent: &tooHigh}), 1e-9)
// exclusive rate below min → clamped to Min
tooLow := -5.0
require.InDelta(t, AffiliateRebateRateMin,
svc.resolveRebateRatePercent(context.Background(), &AffiliateSummary{AffRebateRatePercent: &tooLow}), 1e-9)
}
// TestIsEnabled_NilSettingServiceReturnsDefault verifies that IsEnabled
// safely handles a nil settingService dependency by returning the default
// (off). This protects callers from nil-pointer crashes in misconfigured
// environments.
func TestIsEnabled_NilSettingServiceReturnsDefault(t *testing.T) {
t.Parallel()
svc := &AffiliateService{}
require.False(t, svc.IsEnabled(context.Background()))
require.Equal(t, AffiliateEnabledDefault, svc.IsEnabled(context.Background()))
}
// TestValidateExclusiveRate_BoundaryAndInvalid covers the validator used by
// admin-facing rate setters: nil is always valid (clear), in-range values
// are accepted, NaN/Inf and out-of-range values produce a typed BadRequest.
func TestValidateExclusiveRate_BoundaryAndInvalid(t *testing.T) {
t.Parallel()
require.NoError(t, validateExclusiveRate(nil))
for _, v := range []float64{0, 0.01, 50, 99.99, 100} {
v := v
require.NoError(t, validateExclusiveRate(&v), "value %v should be valid", v)
}
for _, v := range []float64{-0.01, 100.01, -100, 200} {
v := v
require.Error(t, validateExclusiveRate(&v), "value %v should be rejected", v)
}
nan := math.NaN()
require.Error(t, validateExclusiveRate(&nan))
posInf := math.Inf(1)
require.Error(t, validateExclusiveRate(&posInf))
negInf := math.Inf(-1)
require.Error(t, validateExclusiveRate(&negInf))
}
func TestMaskEmail(t *testing.T) {
t.Parallel()
require.Equal(t, "a***@g***.com", maskEmail("alice@gmail.com"))
require.Equal(t, "x***@d***", maskEmail("x@domain"))
require.Equal(t, "", maskEmail(""))
}
func TestIsValidAffiliateCodeFormat(t *testing.T) {
t.Parallel()
// 邀请码格式校验同时服务于:
// 1) 系统自动生成的 12 位随机码(A-Z 去 I/O2-9 去 0/1
// 2) 管理员设置的自定义专属码(如 "VIP2026"、"NEW_USER-1"
// 因此校验放宽到 [A-Z0-9_-]{4,32}(要求调用方先 ToUpper)。
cases := []struct {
name string
in string
want bool
}{
{"valid canonical 12-char", "ABCDEFGHJKLM", true},
{"valid all digits 2-9", "234567892345", true},
{"valid mixed", "A2B3C4D5E6F7", true},
{"valid admin custom short", "VIP1", true},
{"valid admin custom with hyphen", "NEW-USER", true},
{"valid admin custom with underscore", "VIP_2026", true},
{"valid 32-char max", "ABCDEFGHIJKLMNOPQRSTUVWXYZ012345", true},
// Previously-excluded chars (I/O/0/1) are now allowed since admins may use them.
{"letter I now allowed", "IBCDEFGHJKLM", true},
{"letter O now allowed", "OBCDEFGHJKLM", true},
{"digit 0 now allowed", "0BCDEFGHJKLM", true},
{"digit 1 now allowed", "1BCDEFGHJKLM", true},
{"too short (3 chars)", "ABC", false},
{"too long (33 chars)", "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456", false},
{"lowercase rejected (caller must ToUpper first)", "abcdefghjklm", false},
{"empty", "", false},
{"utf8 non-ascii", "ÄÄÄÄÄÄ", false}, // bytes out of charset
{"ascii punctuation .", "ABCDEFGHJK.M", false},
{"whitespace", "ABCDEFGHJK M", false},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
require.Equal(t, tc.want, isValidAffiliateCodeFormat(tc.in))
})
}
}
@@ -0,0 +1,171 @@
package service
import (
"context"
"errors"
"fmt"
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)
var (
ErrAliyunCaptchaVerificationFailed = infraerrors.BadRequest("ALIYUN_CAPTCHA_VERIFICATION_FAILED", "aliyun captcha verification failed")
ErrAliyunCaptchaNotConfigured = infraerrors.ServiceUnavailable("ALIYUN_CAPTCHA_NOT_CONFIGURED", "aliyun captcha not configured")
// ErrCaptchaInvalidCredentials 阿里云验证码凭证无效(仅后台保存校验时返回,公开接口错误码不变)
ErrCaptchaInvalidCredentials = infraerrors.BadRequest("CAPTCHA_INVALID_CREDENTIALS", "invalid aliyun captcha credentials")
)
// AliyunCaptchaCredentials 阿里云验证码 2.0 服务端校验所需的完整凭证
type AliyunCaptchaCredentials struct {
AccessKeyID string
AccessKeySecret string
SceneID string
Endpoint string
}
// AliyunCaptchaVerifyResult VerifyIntelligentCaptcha 的归一化结果
type AliyunCaptchaVerifyResult struct {
VerifyResult bool
VerifyCode string // 阿里云细分结果码,仅用于日志
}
// AliyunCaptchaAPIError 阿里云 OpenAPI 业务错误。
// repository 层负责把 SDK 错误归一化为该类型,service 层不依赖 SDK 包。
type AliyunCaptchaAPIError struct {
Code string
Message string
}
func (e *AliyunCaptchaAPIError) Error() string {
return fmt.Sprintf("aliyun captcha api error: %s: %s", e.Code, e.Message)
}
// AliyunCaptchaVerifier 调用阿里云验证码 2.0 服务端校验的端口
type AliyunCaptchaVerifier interface {
VerifyCaptcha(ctx context.Context, cred AliyunCaptchaCredentials, captchaVerifyParam string) (*AliyunCaptchaVerifyResult, error)
}
const (
// AliyunCaptchaRegionCN 中国内地;AliyunCaptchaRegionSGP 新加坡。
// 该值同时下发给前端 AliyunCaptchaConfig.region,两端必须一致。
AliyunCaptchaRegionCN = "cn"
AliyunCaptchaRegionSGP = "sgp"
aliyunCaptchaEndpointCN = "captcha.cn-shanghai.aliyuncs.com"
aliyunCaptchaEndpointSGP = "captcha.ap-southeast-1.aliyuncs.com"
)
// aliyunCaptchaEndpoint 按后台配置的地域返回服务端接入点,未知值回退中国内地
func aliyunCaptchaEndpoint(region string) string {
if region == AliyunCaptchaRegionSGP {
return aliyunCaptchaEndpointSGP
}
return aliyunCaptchaEndpointCN
}
// normalizeAliyunCaptchaRegion 非法值一律视为中国内地
func normalizeAliyunCaptchaRegion(value string) string {
if value == AliyunCaptchaRegionSGP {
return AliyunCaptchaRegionSGP
}
return AliyunCaptchaRegionCN
}
// aliyunCredentialValidationParam 用于后台保存时探测凭证有效性的假验证参数
const aliyunCredentialValidationParam = "sub2api-credential-validation"
// aliyunInvalidCredentialCodes 表示 AK/SK 本身无效的阿里云错误码;
// 其余错误码(如 param 无效)说明签名已通过、凭证可用。
var aliyunInvalidCredentialCodes = map[string]struct{}{
"InvalidAccessKeyId.NotFound": {},
"InvalidAccessKeyId.Inactive": {},
"SignatureDoesNotMatch": {},
"Forbidden.AccessKeyDisabled": {},
"IncompleteSignature": {},
"InvalidSecurityToken.Expired": {},
}
// AliyunCaptchaService 阿里云验证码 2.0 服务端校验
type AliyunCaptchaService struct {
settingService *SettingService
verifier AliyunCaptchaVerifier
}
func NewAliyunCaptchaService(settingService *SettingService, verifier AliyunCaptchaVerifier) *AliyunCaptchaService {
return &AliyunCaptchaService{settingService: settingService, verifier: verifier}
}
func aliyunCaptchaCredentials(config AliyunCaptchaConfig) (AliyunCaptchaCredentials, bool) {
cred := AliyunCaptchaCredentials{
AccessKeyID: strings.TrimSpace(config.AccessKeyID),
AccessKeySecret: strings.TrimSpace(config.AccessKeySecret),
SceneID: strings.TrimSpace(config.SceneID),
Endpoint: aliyunCaptchaEndpoint(config.Region),
}
if cred.AccessKeyID == "" || cred.AccessKeySecret == "" || cred.SceneID == "" {
return AliyunCaptchaCredentials{}, false
}
return cred, true
}
// VerifyParamWithConfig 校验阿里云验证码 2.0 的 captchaVerifyParam。
// 调用异常时返回错误(fail-closed),与 Turnstile 网络错误行为对称。
func (s *AliyunCaptchaService) VerifyParamWithConfig(ctx context.Context, config AliyunCaptchaConfig, captchaVerifyParam string) error {
if s == nil || s.verifier == nil {
return ErrAliyunCaptchaNotConfigured
}
cred, ok := aliyunCaptchaCredentials(config)
if !ok {
logger.LegacyPrintf("service.aliyun_captcha", "%s", "[AliyunCaptcha] credentials not configured")
return ErrAliyunCaptchaNotConfigured
}
if strings.TrimSpace(captchaVerifyParam) == "" {
logger.LegacyPrintf("service.aliyun_captcha", "%s", "[AliyunCaptcha] captchaVerifyParam is empty")
return ErrAliyunCaptchaVerificationFailed
}
result, err := s.verifier.VerifyCaptcha(ctx, cred, captchaVerifyParam)
if err != nil {
logger.LegacyPrintf("service.aliyun_captcha", "[AliyunCaptcha] verify request failed: %v", err)
return fmt.Errorf("%w: verifier request failed", ErrAliyunCaptchaVerificationFailed)
}
if result == nil || !result.VerifyResult {
if result != nil {
logger.LegacyPrintf("service.aliyun_captcha", "[AliyunCaptcha] rejected, verify code: %s", result.VerifyCode)
}
return ErrAliyunCaptchaVerificationFailed
}
return nil
}
// ValidateCredentials 用假验证参数探测阿里云 AK/SK 是否可用(后台保存设置时调用)。
// 凭证类错误码返回 ErrCaptchaInvalidCredentials;正常响应(包括 param 无效导致的
// VerifyResult=false)说明签名通过、凭证有效;其余错误原样返回给管理员排查。
func (s *AliyunCaptchaService) ValidateCredentials(ctx context.Context, accessKeyID, accessKeySecret, sceneID, region string) error {
if s.verifier == nil {
return ErrAliyunCaptchaNotConfigured
}
cred := AliyunCaptchaCredentials{
AccessKeyID: accessKeyID,
AccessKeySecret: accessKeySecret,
SceneID: sceneID,
Endpoint: aliyunCaptchaEndpoint(region),
}
_, err := s.verifier.VerifyCaptcha(ctx, cred, aliyunCredentialValidationParam)
if err != nil {
var apiErr *AliyunCaptchaAPIError
if errors.As(err, &apiErr) {
if _, invalid := aliyunInvalidCredentialCodes[apiErr.Code]; invalid {
return ErrCaptchaInvalidCredentials
}
}
return fmt.Errorf("validate aliyun captcha credentials: %w", err)
}
return nil
}
@@ -0,0 +1,233 @@
//go:build unit
package service
import (
"context"
"errors"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type aliyunVerifierSpy struct {
called int
lastCred AliyunCaptchaCredentials
lastParam string
result *AliyunCaptchaVerifyResult
err error
}
func (s *aliyunVerifierSpy) VerifyCaptcha(_ context.Context, cred AliyunCaptchaCredentials, param string) (*AliyunCaptchaVerifyResult, error) {
s.called++
s.lastCred = cred
s.lastParam = param
if s.err != nil {
return nil, s.err
}
if s.result != nil {
return s.result, nil
}
return &AliyunCaptchaVerifyResult{VerifyResult: true}, nil
}
func aliyunEnabledSettings() map[string]string {
return map[string]string{
SettingKeyAliyunCaptchaEnabled: "true",
SettingKeyAliyunCaptchaAccessKeyID: "ak-id",
SettingKeyAliyunCaptchaAccessKeySecret: "ak-secret",
SettingKeyAliyunCaptchaSceneID: "scene-1",
SettingKeyAliyunCaptchaPrefix: "prefix-1",
}
}
func aliyunTestConfig() AliyunCaptchaConfig {
return AliyunCaptchaConfig{
Enabled: true,
AccessKeyID: "ak-id",
AccessKeySecret: "ak-secret",
SceneID: "scene-1",
Region: AliyunCaptchaRegionCN,
}
}
func newAliyunAuthServiceForTest(cfg *config.Config, settings map[string]string, aliyunSpy *aliyunVerifierSpy) *AuthService {
settingService := NewSettingService(&settingPublicRepoStub{values: settings}, cfg)
authService := NewAuthService(
nil, // entClient
nil, // userRepo
nil, // redeemRepo
nil, // refreshTokenCache
cfg,
settingService,
nil, // emailService
NewTurnstileService(settingService, &turnstileVerifierSpy{}),
nil, // emailQueueService
nil, // promoService
nil, // defaultSubAssigner
nil, // affiliateService
nil, // userPlatformQuotaRepo
)
authService.SetAliyunCaptchaService(NewAliyunCaptchaService(settingService, aliyunSpy))
return authService
}
func TestAliyunCaptchaServiceVerifyParamDispatch(t *testing.T) {
spy := &aliyunVerifierSpy{}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.VerifyParamWithConfig(context.Background(), aliyunTestConfig(), "captcha-verify-param")
require.NoError(t, err)
require.Equal(t, 1, spy.called)
require.Equal(t, "captcha-verify-param", spy.lastParam)
require.Equal(t, "ak-id", spy.lastCred.AccessKeyID)
require.Equal(t, "scene-1", spy.lastCred.SceneID)
require.Equal(t, "captcha.cn-shanghai.aliyuncs.com", spy.lastCred.Endpoint)
}
func TestAliyunCaptchaServiceSgpEndpoint(t *testing.T) {
spy := &aliyunVerifierSpy{}
svc := NewAliyunCaptchaService(nil, spy)
cfg := aliyunTestConfig()
cfg.Region = AliyunCaptchaRegionSGP
err := svc.VerifyParamWithConfig(context.Background(), cfg, "captcha-verify-param")
require.NoError(t, err)
require.Equal(t, "captcha.ap-southeast-1.aliyuncs.com", spy.lastCred.Endpoint)
}
func TestAliyunCaptchaServiceFailsClosedOnVerifierError(t *testing.T) {
spy := &aliyunVerifierSpy{err: errors.New("network down")}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.VerifyParamWithConfig(context.Background(), aliyunTestConfig(), "captcha-verify-param")
require.ErrorIs(t, err, ErrAliyunCaptchaVerificationFailed)
}
func TestAliyunCaptchaServiceRejectsVerifyResultFalse(t *testing.T) {
spy := &aliyunVerifierSpy{result: &AliyunCaptchaVerifyResult{VerifyResult: false, VerifyCode: "F001"}}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.VerifyParamWithConfig(context.Background(), aliyunTestConfig(), "captcha-verify-param")
require.ErrorIs(t, err, ErrAliyunCaptchaVerificationFailed)
}
func TestAliyunCaptchaServiceRejectsIncompleteCredentials(t *testing.T) {
spy := &aliyunVerifierSpy{}
svc := NewAliyunCaptchaService(nil, spy)
cfg := aliyunTestConfig()
cfg.AccessKeySecret = ""
err := svc.VerifyParamWithConfig(context.Background(), cfg, "captcha-verify-param")
require.ErrorIs(t, err, ErrAliyunCaptchaNotConfigured)
require.Zero(t, spy.called)
}
func TestAliyunCaptchaServiceRejectsEmptyParam(t *testing.T) {
spy := &aliyunVerifierSpy{}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.VerifyParamWithConfig(context.Background(), aliyunTestConfig(), "")
require.ErrorIs(t, err, ErrAliyunCaptchaVerificationFailed)
require.Zero(t, spy.called)
}
func TestAliyunCaptchaServiceValidateCredentials(t *testing.T) {
t.Run("invalid credential code", func(t *testing.T) {
spy := &aliyunVerifierSpy{err: &AliyunCaptchaAPIError{Code: "SignatureDoesNotMatch", Message: "bad sk"}}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.ValidateCredentials(context.Background(), "id", "sk", "scene", "cn")
require.ErrorIs(t, err, ErrCaptchaInvalidCredentials)
})
t.Run("network error surfaces", func(t *testing.T) {
spy := &aliyunVerifierSpy{err: errors.New("timeout")}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.ValidateCredentials(context.Background(), "id", "sk", "scene", "cn")
require.Error(t, err)
require.NotErrorIs(t, err, ErrCaptchaInvalidCredentials)
})
t.Run("verify result false means credentials valid", func(t *testing.T) {
spy := &aliyunVerifierSpy{result: &AliyunCaptchaVerifyResult{VerifyResult: false}}
svc := NewAliyunCaptchaService(nil, spy)
err := svc.ValidateCredentials(context.Background(), "id", "sk", "scene", "sgp")
require.NoError(t, err)
require.Equal(t, "captcha.ap-southeast-1.aliyuncs.com", spy.lastCred.Endpoint)
})
}
func TestAuthServiceVerifyCaptchaDispatchesAliyun(t *testing.T) {
spy := &aliyunVerifierSpy{}
authService := newAliyunAuthServiceForTest(&config.Config{}, aliyunEnabledSettings(), spy)
// 阿里云 captchaVerifyParam 复用 turnstile_token 请求字段
err := authService.VerifyCaptcha(context.Background(), CaptchaProof{TurnstileToken: "captcha-verify-param"}, "127.0.0.1")
require.NoError(t, err)
require.Equal(t, 1, spy.called)
require.Equal(t, "captcha-verify-param", spy.lastParam)
}
func TestAuthServiceVerifyCaptchaRejectsProviderConflict(t *testing.T) {
settings := aliyunEnabledSettings()
settings[SettingKeyTurnstileEnabled] = "true"
settings[SettingKeyTurnstileSecretKey] = "secret"
spy := &aliyunVerifierSpy{}
authService := newAliyunAuthServiceForTest(&config.Config{}, settings, spy)
err := authService.VerifyCaptcha(context.Background(), CaptchaProof{TurnstileToken: "param"}, "127.0.0.1")
require.ErrorIs(t, err, ErrCaptchaProviderConflict)
require.Zero(t, spy.called)
}
func TestAuthServiceVerifyCaptchaRequiredModeWithAliyun(t *testing.T) {
cfg := &config.Config{
Server: config.ServerConfig{Mode: "release"},
Turnstile: config.TurnstileConfig{Required: true},
}
spy := &aliyunVerifierSpy{}
authService := newAliyunAuthServiceForTest(cfg, aliyunEnabledSettings(), spy)
// required 模式 + 阿里云启用且凭证齐全:不误报 NOT_CONFIGURED,正常走阿里云校验
err := authService.VerifyCaptcha(context.Background(), CaptchaProof{TurnstileToken: "captcha-verify-param"}, "127.0.0.1")
require.NoError(t, err)
require.Equal(t, 1, spy.called)
}
func TestAuthServiceVerifyActionCaptchaIfEnabledDispatchesAliyun(t *testing.T) {
spy := &aliyunVerifierSpy{}
authService := newAliyunAuthServiceForTest(&config.Config{}, aliyunEnabledSettings(), spy)
err := authService.VerifyActionCaptchaIfEnabled(context.Background(), CaptchaProof{TurnstileToken: "captcha-verify-param"}, "127.0.0.1")
require.NoError(t, err)
require.Equal(t, 1, spy.called)
require.Equal(t, "captcha-verify-param", spy.lastParam)
}
func TestAuthServiceVerifyActionCaptchaIfEnabledSkipsWhenOnlyTurnstile(t *testing.T) {
spy := &aliyunVerifierSpy{}
authService := newAliyunAuthServiceForTest(&config.Config{}, map[string]string{
SettingKeyTurnstileEnabled: "true",
SettingKeyTurnstileSecretKey: "secret",
}, spy)
// Turnstile 不扩大既有覆盖:扩展入口不拦截
err := authService.VerifyActionCaptchaIfEnabled(context.Background(), CaptchaProof{}, "127.0.0.1")
require.NoError(t, err)
require.Zero(t, spy.called)
}
+85
View File
@@ -0,0 +1,85 @@
package service
import (
"context"
"time"
"github.com/Wei-Shaw/sub2api/internal/domain"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
)
const (
AnnouncementStatusDraft = domain.AnnouncementStatusDraft
AnnouncementStatusActive = domain.AnnouncementStatusActive
AnnouncementStatusArchived = domain.AnnouncementStatusArchived
)
const (
AnnouncementNotifyModeSilent = domain.AnnouncementNotifyModeSilent
AnnouncementNotifyModePopup = domain.AnnouncementNotifyModePopup
)
const (
AnnouncementConditionTypeSubscription = domain.AnnouncementConditionTypeSubscription
AnnouncementConditionTypeBalance = domain.AnnouncementConditionTypeBalance
)
const (
AnnouncementOperatorIn = domain.AnnouncementOperatorIn
AnnouncementOperatorGT = domain.AnnouncementOperatorGT
AnnouncementOperatorGTE = domain.AnnouncementOperatorGTE
AnnouncementOperatorLT = domain.AnnouncementOperatorLT
AnnouncementOperatorLTE = domain.AnnouncementOperatorLTE
AnnouncementOperatorEQ = domain.AnnouncementOperatorEQ
)
var (
ErrAnnouncementNotFound = domain.ErrAnnouncementNotFound
ErrAnnouncementInvalidTarget = domain.ErrAnnouncementInvalidTarget
ErrAnnouncementNilInput = infraerrors.BadRequest("ANNOUNCEMENT_INPUT_REQUIRED", "announcement input is required")
ErrAnnouncementInvalidTitle = infraerrors.BadRequest("ANNOUNCEMENT_TITLE_INVALID", "announcement title is invalid")
ErrAnnouncementContentRequired = infraerrors.BadRequest(
"ANNOUNCEMENT_CONTENT_REQUIRED",
"announcement content is required",
)
ErrAnnouncementInvalidStatus = infraerrors.BadRequest("ANNOUNCEMENT_STATUS_INVALID", "announcement status is invalid")
ErrAnnouncementInvalidNotifyMode = infraerrors.BadRequest(
"ANNOUNCEMENT_NOTIFY_MODE_INVALID",
"announcement notify_mode is invalid",
)
ErrAnnouncementInvalidSchedule = infraerrors.BadRequest(
"ANNOUNCEMENT_TIME_RANGE_INVALID",
"starts_at must be before ends_at",
)
)
type AnnouncementTargeting = domain.AnnouncementTargeting
type AnnouncementConditionGroup = domain.AnnouncementConditionGroup
type AnnouncementCondition = domain.AnnouncementCondition
type Announcement = domain.Announcement
type AnnouncementListFilters struct {
Status string
Search string
}
type AnnouncementRepository interface {
Create(ctx context.Context, a *Announcement) error
GetByID(ctx context.Context, id int64) (*Announcement, error)
Update(ctx context.Context, a *Announcement) error
Delete(ctx context.Context, id int64) error
List(ctx context.Context, params pagination.PaginationParams, filters AnnouncementListFilters) ([]Announcement, *pagination.PaginationResult, error)
ListActive(ctx context.Context, now time.Time) ([]Announcement, error)
}
type AnnouncementReadRepository interface {
MarkRead(ctx context.Context, announcementID, userID int64, readAt time.Time) error
GetReadMapByUser(ctx context.Context, userID int64, announcementIDs []int64) (map[int64]time.Time, error)
GetReadMapByUsers(ctx context.Context, announcementID int64, userIDs []int64) (map[int64]time.Time, error)
CountByAnnouncementID(ctx context.Context, announcementID int64) (int64, error)
}
@@ -0,0 +1,406 @@
package service
import (
"context"
"fmt"
"sort"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
)
type AnnouncementService struct {
announcementRepo AnnouncementRepository
readRepo AnnouncementReadRepository
userRepo UserRepository
userSubRepo UserSubscriptionRepository
}
func NewAnnouncementService(
announcementRepo AnnouncementRepository,
readRepo AnnouncementReadRepository,
userRepo UserRepository,
userSubRepo UserSubscriptionRepository,
) *AnnouncementService {
return &AnnouncementService{
announcementRepo: announcementRepo,
readRepo: readRepo,
userRepo: userRepo,
userSubRepo: userSubRepo,
}
}
type CreateAnnouncementInput struct {
Title string
Content string
Status string
NotifyMode string
Targeting AnnouncementTargeting
StartsAt *time.Time
EndsAt *time.Time
ActorID *int64 // 管理员用户ID
}
type UpdateAnnouncementInput struct {
Title *string
Content *string
Status *string
NotifyMode *string
Targeting *AnnouncementTargeting
StartsAt **time.Time
EndsAt **time.Time
ActorID *int64 // 管理员用户ID
}
type UserAnnouncement struct {
Announcement Announcement
ReadAt *time.Time
}
type AnnouncementUserReadStatus struct {
UserID int64 `json:"user_id"`
Email string `json:"email"`
Username string `json:"username"`
Balance float64 `json:"balance"`
Eligible bool `json:"eligible"`
ReadAt *time.Time `json:"read_at,omitempty"`
}
func (s *AnnouncementService) Create(ctx context.Context, input *CreateAnnouncementInput) (*Announcement, error) {
if input == nil {
return nil, ErrAnnouncementNilInput
}
title := strings.TrimSpace(input.Title)
content := strings.TrimSpace(input.Content)
if title == "" || len(title) > 200 {
return nil, ErrAnnouncementInvalidTitle
}
if content == "" {
return nil, ErrAnnouncementContentRequired
}
status := strings.TrimSpace(input.Status)
if status == "" {
status = AnnouncementStatusDraft
}
if !isValidAnnouncementStatus(status) {
return nil, ErrAnnouncementInvalidStatus
}
targeting, err := domain.AnnouncementTargeting(input.Targeting).NormalizeAndValidate()
if err != nil {
return nil, err
}
notifyMode := strings.TrimSpace(input.NotifyMode)
if notifyMode == "" {
notifyMode = AnnouncementNotifyModeSilent
}
if !isValidAnnouncementNotifyMode(notifyMode) {
return nil, ErrAnnouncementInvalidNotifyMode
}
if input.StartsAt != nil && input.EndsAt != nil {
if !input.StartsAt.Before(*input.EndsAt) {
return nil, ErrAnnouncementInvalidSchedule
}
}
a := &Announcement{
Title: title,
Content: content,
Status: status,
NotifyMode: notifyMode,
Targeting: targeting,
StartsAt: input.StartsAt,
EndsAt: input.EndsAt,
}
if input.ActorID != nil && *input.ActorID > 0 {
a.CreatedBy = input.ActorID
a.UpdatedBy = input.ActorID
}
if err := s.announcementRepo.Create(ctx, a); err != nil {
return nil, fmt.Errorf("create announcement: %w", err)
}
return a, nil
}
func (s *AnnouncementService) Update(ctx context.Context, id int64, input *UpdateAnnouncementInput) (*Announcement, error) {
if input == nil {
return nil, ErrAnnouncementNilInput
}
a, err := s.announcementRepo.GetByID(ctx, id)
if err != nil {
return nil, err
}
if input.Title != nil {
title := strings.TrimSpace(*input.Title)
if title == "" || len(title) > 200 {
return nil, ErrAnnouncementInvalidTitle
}
a.Title = title
}
if input.Content != nil {
content := strings.TrimSpace(*input.Content)
if content == "" {
return nil, ErrAnnouncementContentRequired
}
a.Content = content
}
if input.Status != nil {
status := strings.TrimSpace(*input.Status)
if !isValidAnnouncementStatus(status) {
return nil, ErrAnnouncementInvalidStatus
}
a.Status = status
}
if input.NotifyMode != nil {
notifyMode := strings.TrimSpace(*input.NotifyMode)
if !isValidAnnouncementNotifyMode(notifyMode) {
return nil, ErrAnnouncementInvalidNotifyMode
}
a.NotifyMode = notifyMode
}
if input.Targeting != nil {
targeting, err := domain.AnnouncementTargeting(*input.Targeting).NormalizeAndValidate()
if err != nil {
return nil, err
}
a.Targeting = targeting
}
if input.StartsAt != nil {
a.StartsAt = *input.StartsAt
}
if input.EndsAt != nil {
a.EndsAt = *input.EndsAt
}
if a.StartsAt != nil && a.EndsAt != nil {
if !a.StartsAt.Before(*a.EndsAt) {
return nil, ErrAnnouncementInvalidSchedule
}
}
if input.ActorID != nil && *input.ActorID > 0 {
a.UpdatedBy = input.ActorID
}
if err := s.announcementRepo.Update(ctx, a); err != nil {
return nil, fmt.Errorf("update announcement: %w", err)
}
return a, nil
}
func (s *AnnouncementService) Delete(ctx context.Context, id int64) error {
if err := s.announcementRepo.Delete(ctx, id); err != nil {
return fmt.Errorf("delete announcement: %w", err)
}
return nil
}
func (s *AnnouncementService) GetByID(ctx context.Context, id int64) (*Announcement, error) {
return s.announcementRepo.GetByID(ctx, id)
}
func (s *AnnouncementService) List(ctx context.Context, params pagination.PaginationParams, filters AnnouncementListFilters) ([]Announcement, *pagination.PaginationResult, error) {
return s.announcementRepo.List(ctx, params, filters)
}
func (s *AnnouncementService) ListForUser(ctx context.Context, userID int64, unreadOnly bool) ([]UserAnnouncement, error) {
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user: %w", err)
}
activeSubs, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("list active subscriptions: %w", err)
}
activeGroupIDs := make(map[int64]struct{}, len(activeSubs))
for i := range activeSubs {
activeGroupIDs[activeSubs[i].GroupID] = struct{}{}
}
now := time.Now()
anns, err := s.announcementRepo.ListActive(ctx, now)
if err != nil {
return nil, fmt.Errorf("list active announcements: %w", err)
}
visible := make([]Announcement, 0, len(anns))
ids := make([]int64, 0, len(anns))
for i := range anns {
a := anns[i]
if !a.IsActiveAt(now) {
continue
}
if !a.Targeting.Matches(user.Balance, activeGroupIDs) {
continue
}
visible = append(visible, a)
ids = append(ids, a.ID)
}
if len(visible) == 0 {
return []UserAnnouncement{}, nil
}
readMap, err := s.readRepo.GetReadMapByUser(ctx, userID, ids)
if err != nil {
return nil, fmt.Errorf("get read map: %w", err)
}
out := make([]UserAnnouncement, 0, len(visible))
for i := range visible {
a := visible[i]
readAt, ok := readMap[a.ID]
if unreadOnly && ok {
continue
}
var ptr *time.Time
if ok {
t := readAt
ptr = &t
}
out = append(out, UserAnnouncement{
Announcement: a,
ReadAt: ptr,
})
}
// 未读优先、同状态按创建时间倒序
sort.Slice(out, func(i, j int) bool {
ai, aj := out[i], out[j]
if (ai.ReadAt == nil) != (aj.ReadAt == nil) {
return ai.ReadAt == nil
}
return ai.Announcement.ID > aj.Announcement.ID
})
return out, nil
}
func (s *AnnouncementService) MarkRead(ctx context.Context, userID, announcementID int64) error {
// 安全:仅允许标记当前用户“可见”的公告
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return fmt.Errorf("get user: %w", err)
}
a, err := s.announcementRepo.GetByID(ctx, announcementID)
if err != nil {
return err
}
now := time.Now()
if !a.IsActiveAt(now) {
return ErrAnnouncementNotFound
}
activeSubs, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
if err != nil {
return fmt.Errorf("list active subscriptions: %w", err)
}
activeGroupIDs := make(map[int64]struct{}, len(activeSubs))
for i := range activeSubs {
activeGroupIDs[activeSubs[i].GroupID] = struct{}{}
}
if !a.Targeting.Matches(user.Balance, activeGroupIDs) {
return ErrAnnouncementNotFound
}
if err := s.readRepo.MarkRead(ctx, announcementID, userID, now); err != nil {
return fmt.Errorf("mark read: %w", err)
}
return nil
}
func (s *AnnouncementService) ListUserReadStatus(
ctx context.Context,
announcementID int64,
params pagination.PaginationParams,
search string,
) ([]AnnouncementUserReadStatus, *pagination.PaginationResult, error) {
ann, err := s.announcementRepo.GetByID(ctx, announcementID)
if err != nil {
return nil, nil, err
}
filters := UserListFilters{
Search: strings.TrimSpace(search),
}
users, page, err := s.userRepo.ListWithFilters(ctx, params, filters)
if err != nil {
return nil, nil, fmt.Errorf("list users: %w", err)
}
userIDs := make([]int64, 0, len(users))
for i := range users {
userIDs = append(userIDs, users[i].ID)
}
readMap, err := s.readRepo.GetReadMapByUsers(ctx, announcementID, userIDs)
if err != nil {
return nil, nil, fmt.Errorf("get read map: %w", err)
}
out := make([]AnnouncementUserReadStatus, 0, len(users))
for i := range users {
u := users[i]
subs, err := s.userSubRepo.ListActiveByUserID(ctx, u.ID)
if err != nil {
return nil, nil, fmt.Errorf("list active subscriptions: %w", err)
}
activeGroupIDs := make(map[int64]struct{}, len(subs))
for j := range subs {
activeGroupIDs[subs[j].GroupID] = struct{}{}
}
readAt, ok := readMap[u.ID]
var ptr *time.Time
if ok {
t := readAt
ptr = &t
}
out = append(out, AnnouncementUserReadStatus{
UserID: u.ID,
Email: u.Email,
Username: u.Username,
Balance: u.Balance,
Eligible: domain.AnnouncementTargeting(ann.Targeting).Matches(u.Balance, activeGroupIDs),
ReadAt: ptr,
})
}
return out, page, nil
}
func isValidAnnouncementStatus(status string) bool {
switch status {
case AnnouncementStatusDraft, AnnouncementStatusActive, AnnouncementStatusArchived:
return true
default:
return false
}
}
func isValidAnnouncementNotifyMode(mode string) bool {
switch mode {
case AnnouncementNotifyModeSilent, AnnouncementNotifyModePopup:
return true
default:
return false
}
}
@@ -0,0 +1,81 @@
package service
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
type announcementRepoStub struct {
item *Announcement
}
func (s *announcementRepoStub) Create(_ context.Context, a *Announcement) error {
s.item = a
return nil
}
func (s *announcementRepoStub) GetByID(_ context.Context, _ int64) (*Announcement, error) {
if s.item == nil {
return nil, ErrAnnouncementNotFound
}
return s.item, nil
}
func (s *announcementRepoStub) Update(_ context.Context, a *Announcement) error {
s.item = a
return nil
}
func (*announcementRepoStub) Delete(context.Context, int64) error {
return nil
}
func (*announcementRepoStub) List(context.Context, pagination.PaginationParams, AnnouncementListFilters) ([]Announcement, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (*announcementRepoStub) ListActive(context.Context, time.Time) ([]Announcement, error) {
return nil, nil
}
func TestAnnouncementServiceCreateRejectsEqualStartEndTimes(t *testing.T) {
repo := &announcementRepoStub{}
svc := NewAnnouncementService(repo, nil, nil, nil)
now := time.Unix(1776790020, 0)
_, err := svc.Create(context.Background(), &CreateAnnouncementInput{
Title: "公告",
Content: "内容",
Status: AnnouncementStatusActive,
NotifyMode: AnnouncementNotifyModePopup,
StartsAt: &now,
EndsAt: &now,
})
require.ErrorIs(t, err, ErrAnnouncementInvalidSchedule)
}
func TestAnnouncementServiceUpdateRejectsEqualStartEndTimes(t *testing.T) {
repo := &announcementRepoStub{
item: &Announcement{
ID: 1,
Title: "公告",
Content: "内容",
Status: AnnouncementStatusActive,
NotifyMode: AnnouncementNotifyModePopup,
},
}
svc := NewAnnouncementService(repo, nil, nil, nil)
now := time.Unix(1776790020, 0)
startsAt := &now
endsAt := &now
_, err := svc.Update(context.Background(), 1, &UpdateAnnouncementInput{
StartsAt: &startsAt,
EndsAt: &endsAt,
})
require.ErrorIs(t, err, ErrAnnouncementInvalidSchedule)
}
@@ -0,0 +1,66 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAnnouncementTargeting_Matches_EmptyMatchesAll(t *testing.T) {
var targeting AnnouncementTargeting
require.True(t, targeting.Matches(0, nil))
require.True(t, targeting.Matches(123.45, map[int64]struct{}{1: {}}))
}
func TestAnnouncementTargeting_NormalizeAndValidate_RejectsEmptyGroup(t *testing.T) {
targeting := AnnouncementTargeting{
AnyOf: []AnnouncementConditionGroup{
{AllOf: nil},
},
}
_, err := targeting.NormalizeAndValidate()
require.Error(t, err)
require.ErrorIs(t, err, ErrAnnouncementInvalidTarget)
}
func TestAnnouncementTargeting_NormalizeAndValidate_RejectsInvalidCondition(t *testing.T) {
targeting := AnnouncementTargeting{
AnyOf: []AnnouncementConditionGroup{
{
AllOf: []AnnouncementCondition{
{Type: "balance", Operator: "between", Value: 10},
},
},
},
}
_, err := targeting.NormalizeAndValidate()
require.Error(t, err)
require.ErrorIs(t, err, ErrAnnouncementInvalidTarget)
}
func TestAnnouncementTargeting_Matches_AndOrSemantics(t *testing.T) {
targeting := AnnouncementTargeting{
AnyOf: []AnnouncementConditionGroup{
{
AllOf: []AnnouncementCondition{
{Type: AnnouncementConditionTypeBalance, Operator: AnnouncementOperatorGTE, Value: 100},
{Type: AnnouncementConditionTypeSubscription, Operator: AnnouncementOperatorIn, GroupIDs: []int64{10}},
},
},
{
AllOf: []AnnouncementCondition{
{Type: AnnouncementConditionTypeBalance, Operator: AnnouncementOperatorLT, Value: 5},
},
},
},
}
// 命中第 2 组(balance < 5
require.True(t, targeting.Matches(4.99, nil))
require.False(t, targeting.Matches(5, nil))
// 命中第 1 组(balance >= 100 AND 订阅 in [10]
require.False(t, targeting.Matches(100, map[int64]struct{}{}))
require.False(t, targeting.Matches(99.9, map[int64]struct{}{10: {}}))
require.True(t, targeting.Matches(100, map[int64]struct{}{10: {}}))
}
@@ -0,0 +1,42 @@
package service
import (
"net/http"
"strings"
)
const (
anthropicAPIKeyAuthSchemeExtraKey = "anthropic_apikey_auth_scheme"
AnthropicAPIKeyAuthSchemeXAPIKey = "x_api_key"
AnthropicAPIKeyAuthSchemeAuthorizationBearer = "authorization_bearer"
)
// GetAnthropicAPIKeyAuthScheme returns the upstream authentication scheme for
// Anthropic API-key accounts. Missing or invalid values keep the historical
// x-api-key behavior. CN providers using their native Anthropic endpoints
// (api_protocol=anthropic) share the same override knob — Kimi/DeepSeek default
// to x-api-key, Zhipu can opt into Authorization: Bearer.
func (a *Account) GetAnthropicAPIKeyAuthScheme() string {
if a == nil || a.Type != AccountTypeAPIKey {
return AnthropicAPIKeyAuthSchemeXAPIKey
}
if a.Platform != PlatformAnthropic && !a.IsCNProvider() {
return AnthropicAPIKeyAuthSchemeXAPIKey
}
switch strings.TrimSpace(a.GetExtraString(anthropicAPIKeyAuthSchemeExtraKey)) {
case AnthropicAPIKeyAuthSchemeAuthorizationBearer:
return AnthropicAPIKeyAuthSchemeAuthorizationBearer
default:
return AnthropicAPIKeyAuthSchemeXAPIKey
}
}
func setAnthropicAPIKeyAuthHeader(header http.Header, account *Account, token string) {
if account.GetAnthropicAPIKeyAuthScheme() == AnthropicAPIKeyAuthSchemeAuthorizationBearer {
header.Set("Authorization", "Bearer "+token)
return
}
header.Set("x-api-key", token)
}
@@ -0,0 +1,90 @@
package service
import (
"encoding/json"
"strings"
"time"
"github.com/tidwall/gjson"
)
// Anthropic 会话 Fallback 相关常量
const (
// anthropicSessionTTLSeconds Anthropic 会话缓存 TTL5 分钟)
anthropicSessionTTLSeconds = 300
// anthropicDigestSessionKeyPrefix Anthropic 摘要 fallback 会话 key 前缀
anthropicDigestSessionKeyPrefix = "anthropic:digest:"
)
// AnthropicSessionTTL 返回 Anthropic 会话缓存 TTL
func AnthropicSessionTTL() time.Duration {
return anthropicSessionTTLSeconds * time.Second
}
// BuildAnthropicDigestChain 根据 Anthropic 请求生成摘要链
// 格式: s:<hash>-u:<hash>-a:<hash>-u:<hash>-...
// s = system, u = user, a = assistant
func BuildAnthropicDigestChain(parsed *ParsedRequest) string {
if parsed == nil {
return ""
}
var parts []string
if systemRaw := parsed.SystemRaw(); len(systemRaw) > 0 && string(systemRaw) != "null" {
parts = append(parts, "s:"+shortHash(canonicalAnthropicDigestJSON(systemRaw)))
}
messages := parsed.MessagesRaw()
if len(messages) > 0 {
gjson.ParseBytes(messages).ForEach(func(_, msg gjson.Result) bool {
prefix := rolePrefix(msg.Get("role").String())
content := msg.Get("content")
parts = append(parts, prefix+":"+shortHash(canonicalAnthropicDigestJSON([]byte(content.Raw))))
return true
})
}
return strings.Join(parts, "-")
}
// canonicalAnthropicDigestJSON 保持 digest 对 JSON key 顺序和空白不敏感。
func canonicalAnthropicDigestJSON(raw []byte) []byte {
if len(raw) == 0 {
return raw
}
var value any
if err := json.Unmarshal(raw, &value); err != nil {
return raw
}
canonical, err := json.Marshal(value)
if err != nil {
return raw
}
return canonical
}
// rolePrefix 将 Anthropic 的 role 映射为单字符前缀
func rolePrefix(role string) string {
switch role {
case "assistant":
return "a"
default:
return "u"
}
}
// GenerateAnthropicDigestSessionKey 生成 Anthropic 摘要 fallback 的 sessionKey
// 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey
func GenerateAnthropicDigestSessionKey(prefixHash, uuid string) string {
prefix := prefixHash
if len(prefixHash) >= 8 {
prefix = prefixHash[:8]
}
uuidPart := uuid
if len(uuid) >= 8 {
uuidPart = uuid[:8]
}
return anthropicDigestSessionKeyPrefix + prefix + ":" + uuidPart
}

Some files were not shown because too many files have changed in this diff Show More