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
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:
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)
|
||||
})
|
||||
}
|
||||
@@ -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 校验并规范化单个覆写条目,保存路径(Normalize,err → 400)
|
||||
// 与应用路径(resolveHeaderOverrides,err → 跳过)共用同一套规则,避免两处校验漂移。
|
||||
// 名和值均为空表示空占位行,返回 ("", "", 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%")
|
||||
}
|
||||
@@ -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-#5641:compact_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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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]*rate,nil 表示删除该分组的专属倍率
|
||||
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: ¬es,
|
||||
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/O,2-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)
|
||||
}
|
||||
@@ -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 会话缓存 TTL(5 分钟)
|
||||
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
Reference in New Issue
Block a user