347 lines
13 KiB
Go
347 lines
13 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"net/http"
|
|||
|
|
"net/http/httptest"
|
|||
|
|
"testing"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/imroc/req/v3"
|
|||
|
|
"github.com/stretchr/testify/require"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
func TestFetchChatGPTSubscriptionExpiresAt(t *testing.T) {
|
|||
|
|
const wantExpiresAt = "2026-06-10T02:52:15Z"
|
|||
|
|
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|||
|
|
require.Equal(t, "/backend-api/subscriptions", r.URL.Path)
|
|||
|
|
require.Equal(t, "acc_123", r.URL.Query().Get("account_id"))
|
|||
|
|
require.Equal(t, "Bearer access-token", r.Header.Get("Authorization"))
|
|||
|
|
|
|||
|
|
w.Header().Set("Content-Type", "application/json")
|
|||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|||
|
|
"plan_type": "plus",
|
|||
|
|
"active_until": wantExpiresAt,
|
|||
|
|
"will_renew": true,
|
|||
|
|
"id": "sub_123",
|
|||
|
|
})
|
|||
|
|
}))
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
oldURL := chatGPTSubscriptionsURL
|
|||
|
|
chatGPTSubscriptionsURL = server.URL + "/backend-api/subscriptions"
|
|||
|
|
t.Cleanup(func() { chatGPTSubscriptionsURL = oldURL })
|
|||
|
|
|
|||
|
|
got := fetchChatGPTSubscriptionExpiresAt(context.Background(), func(proxyURL string) (*req.Client, error) {
|
|||
|
|
return req.C().SetTimeout(5 * time.Second), nil
|
|||
|
|
}, "access-token", "", "acc_123")
|
|||
|
|
|
|||
|
|
require.Equal(t, wantExpiresAt, got)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestFetchChatGPTAccountInfo_SkipsExpiredWorkspaceCandidate(t *testing.T) {
|
|||
|
|
expiredAt := time.Now().Add(-24 * time.Hour).UTC().Format(time.RFC3339)
|
|||
|
|
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|||
|
|
require.Equal(t, "/backend-api/accounts/check/v4-2023-04-27", r.URL.Path)
|
|||
|
|
require.Equal(t, "Bearer access-token", r.Header.Get("Authorization"))
|
|||
|
|
|
|||
|
|
w.Header().Set("Content-Type", "application/json")
|
|||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
"org-expired-workspace": map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"plan_type": "self_serve_business_usage_based",
|
|||
|
|
"is_default": true,
|
|||
|
|
},
|
|||
|
|
"entitlement": map[string]any{
|
|||
|
|
"expires_at": expiredAt,
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
"personal-account": map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"plan_type": "free",
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
}))
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
oldURL := chatGPTAccountsCheckURL
|
|||
|
|
chatGPTAccountsCheckURL = server.URL + "/backend-api/accounts/check/v4-2023-04-27"
|
|||
|
|
t.Cleanup(func() { chatGPTAccountsCheckURL = oldURL })
|
|||
|
|
|
|||
|
|
got := fetchChatGPTAccountInfo(context.Background(), func(proxyURL string) (*req.Client, error) {
|
|||
|
|
return req.C().SetTimeout(5 * time.Second), nil
|
|||
|
|
}, "access-token", "", "org-expired-workspace")
|
|||
|
|
|
|||
|
|
require.NotNil(t, got)
|
|||
|
|
require.Equal(t, "free", got.PlanType)
|
|||
|
|
require.Empty(t, got.SubscriptionExpiresAt)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestFetchChatGPTAccountInfo_SkipsDeactivatedWorkspaceCandidate(t *testing.T) {
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|||
|
|
require.Equal(t, "/backend-api/accounts/check/v4-2023-04-27", r.URL.Path)
|
|||
|
|
|
|||
|
|
w.Header().Set("Content-Type", "application/json")
|
|||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
"org-deactivated-workspace": map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"plan_type": "self_serve_business_usage_based",
|
|||
|
|
"is_default": true,
|
|||
|
|
"is_deactivated": true,
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
"personal-account": map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"plan_type": "pro",
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
}))
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
oldURL := chatGPTAccountsCheckURL
|
|||
|
|
chatGPTAccountsCheckURL = server.URL + "/backend-api/accounts/check/v4-2023-04-27"
|
|||
|
|
t.Cleanup(func() { chatGPTAccountsCheckURL = oldURL })
|
|||
|
|
|
|||
|
|
got := fetchChatGPTAccountInfo(context.Background(), func(proxyURL string) (*req.Client, error) {
|
|||
|
|
return req.C().SetTimeout(5 * time.Second), nil
|
|||
|
|
}, "access-token", "", "org-deactivated-workspace")
|
|||
|
|
|
|||
|
|
require.NotNil(t, got)
|
|||
|
|
require.Equal(t, "pro", got.PlanType)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestShouldApplyChatGPTAccountInfoPlanType(t *testing.T) {
|
|||
|
|
require.False(t, shouldApplyChatGPTAccountInfoPlanType("pro", "self_serve_business_usage_based"))
|
|||
|
|
require.False(t, shouldApplyChatGPTAccountInfoPlanType("free", "team"))
|
|||
|
|
require.False(t, shouldApplyChatGPTAccountInfoPlanType("", ""))
|
|||
|
|
require.True(t, shouldApplyChatGPTAccountInfoPlanType("", "pro"))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestChatGPTAccountInfoBelongsToTokenAccount(t *testing.T) {
|
|||
|
|
require.False(t, chatGPTAccountInfoBelongsToTokenAccount(
|
|||
|
|
&OpenAITokenInfo{ChatGPTAccountID: "personal-a"}, &ChatGPTAccountInfo{AccountID: "workspace-b"}))
|
|||
|
|
require.True(t, chatGPTAccountInfoBelongsToTokenAccount(
|
|||
|
|
&OpenAITokenInfo{ChatGPTAccountID: "personal-a"}, &ChatGPTAccountInfo{AccountID: "PERSONAL-A"}))
|
|||
|
|
// 任一侧缺 ID 时无法区分，保持既有行为（采用 accounts/check 的值）。
|
|||
|
|
require.True(t, chatGPTAccountInfoBelongsToTokenAccount(
|
|||
|
|
&OpenAITokenInfo{}, &ChatGPTAccountInfo{AccountID: "workspace-b"}))
|
|||
|
|
require.True(t, chatGPTAccountInfoBelongsToTokenAccount(
|
|||
|
|
&OpenAITokenInfo{ChatGPTAccountID: "personal-a"}, &ChatGPTAccountInfo{}))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// accounts 的 map key 可能是 "default" 这类别名，account.account_id 才是账号标识。
|
|||
|
|
func TestFetchChatGPTAccountInfo_ReportsAccountID(t *testing.T) {
|
|||
|
|
futureAt := time.Now().Add(720 * time.Hour).UTC().Format(time.RFC3339)
|
|||
|
|
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|||
|
|
w.Header().Set("Content-Type", "application/json")
|
|||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
"default": map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"account_id": "personal-account-a",
|
|||
|
|
"plan_type": "plus",
|
|||
|
|
"is_default": true,
|
|||
|
|
},
|
|||
|
|
"entitlement": map[string]any{"expires_at": futureAt},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
}))
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
oldURL := chatGPTAccountsCheckURL
|
|||
|
|
chatGPTAccountsCheckURL = server.URL + "/backend-api/accounts/check/v4-2023-04-27"
|
|||
|
|
t.Cleanup(func() { chatGPTAccountsCheckURL = oldURL })
|
|||
|
|
|
|||
|
|
got := fetchChatGPTAccountInfo(context.Background(), newTestPrivacyClientFactory(), "access-token", "", "")
|
|||
|
|
require.NotNil(t, got)
|
|||
|
|
require.Equal(t, "plus", got.PlanType)
|
|||
|
|
require.Equal(t, "personal-account-a", got.AccountID, "应优先取 account.account_id 而不是 map key")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// issue #5459：poid 指向的默认 Personal workspace 与 chatgpt_account_id 是两个不同的
|
|||
|
|
// 标识时，accounts/check 返回的是 workspace 的 entitlement.expires_at。plan_type 那侧
|
|||
|
|
// 已被 shouldApplyChatGPTAccountInfoPlanType 挡住（保留 JWT 里的个人套餐），到期时间
|
|||
|
|
// 这侧原先无条件覆盖，于是显示成「个人 Pro + workspace 到期时间」。
|
|||
|
|
func TestEnrichTokenInfo_WorkspaceEntitlementDoesNotOverridePersonalSubscription(t *testing.T) {
|
|||
|
|
const (
|
|||
|
|
personalAccountID = "personal-account-a"
|
|||
|
|
workspaceAccountID = "personal-workspace-b"
|
|||
|
|
personalActiveUntil = "2027-03-01T00:00:00Z"
|
|||
|
|
)
|
|||
|
|
workspaceExpiresAt := time.Now().Add(720 * time.Hour).UTC().Format(time.RFC3339)
|
|||
|
|
|
|||
|
|
subscriptionCalls := 0
|
|||
|
|
server := newChatGPTBackendTestServer(t, chatGPTBackendTestServerConfig{
|
|||
|
|
accountsCheck: map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
workspaceAccountID: map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"account_id": workspaceAccountID,
|
|||
|
|
"plan_type": "pro",
|
|||
|
|
"is_default": true,
|
|||
|
|
},
|
|||
|
|
"entitlement": map[string]any{"expires_at": workspaceExpiresAt},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
onSubscription: func(accountID string) map[string]any {
|
|||
|
|
subscriptionCalls++
|
|||
|
|
require.Equal(t, personalAccountID, accountID,
|
|||
|
|
"必须用个人 chatgpt_account_id 查订阅，而不是 poid workspace")
|
|||
|
|
return map[string]any{"plan_type": "pro", "active_until": personalActiveUntil, "will_renew": true}
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
tokenInfo := &OpenAITokenInfo{
|
|||
|
|
AccessToken: "access-token",
|
|||
|
|
ChatGPTAccountID: personalAccountID,
|
|||
|
|
OrganizationID: workspaceAccountID,
|
|||
|
|
PlanType: "pro", // 来自 id_token 的个人套餐
|
|||
|
|
}
|
|||
|
|
svc := &OpenAIOAuthService{privacyClientFactory: newTestPrivacyClientFactory()}
|
|||
|
|
svc.enrichTokenInfo(context.Background(), tokenInfo, "")
|
|||
|
|
|
|||
|
|
require.Equal(t, "pro", tokenInfo.PlanType)
|
|||
|
|
require.Equal(t, personalActiveUntil, tokenInfo.SubscriptionExpiresAt,
|
|||
|
|
"到期时间必须来自个人订阅 active_until，不能是 workspace 的 entitlement.expires_at")
|
|||
|
|
require.NotEqual(t, workspaceExpiresAt, tokenInfo.SubscriptionExpiresAt)
|
|||
|
|
require.Equal(t, 1, subscriptionCalls)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 单个人账号（poid == chatgpt_account_id）是绝大多数情况，行为必须保持不变：
|
|||
|
|
// 直接用 accounts/check 的 entitlement，不额外打订阅端点。
|
|||
|
|
func TestEnrichTokenInfo_KeepsEntitlementWhenAccountMatches(t *testing.T) {
|
|||
|
|
const personalAccountID = "personal-account-a"
|
|||
|
|
entitlementExpiresAt := time.Now().Add(720 * time.Hour).UTC().Format(time.RFC3339)
|
|||
|
|
|
|||
|
|
subscriptionCalls := 0
|
|||
|
|
server := newChatGPTBackendTestServer(t, chatGPTBackendTestServerConfig{
|
|||
|
|
accountsCheck: map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
personalAccountID: map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"account_id": personalAccountID,
|
|||
|
|
"plan_type": "plus",
|
|||
|
|
"is_default": true,
|
|||
|
|
},
|
|||
|
|
"entitlement": map[string]any{"expires_at": entitlementExpiresAt},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
onSubscription: func(string) map[string]any {
|
|||
|
|
subscriptionCalls++
|
|||
|
|
return map[string]any{}
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
tokenInfo := &OpenAITokenInfo{
|
|||
|
|
AccessToken: "access-token",
|
|||
|
|
ChatGPTAccountID: personalAccountID,
|
|||
|
|
OrganizationID: personalAccountID,
|
|||
|
|
PlanType: "plus",
|
|||
|
|
}
|
|||
|
|
svc := &OpenAIOAuthService{privacyClientFactory: newTestPrivacyClientFactory()}
|
|||
|
|
svc.enrichTokenInfo(context.Background(), tokenInfo, "")
|
|||
|
|
|
|||
|
|
require.Equal(t, entitlementExpiresAt, tokenInfo.SubscriptionExpiresAt)
|
|||
|
|
require.Zero(t, subscriptionCalls, "账号一致时不应额外请求订阅端点")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 反向不变式：套餐本身就取自 accounts/check（JWT 没有 plan_type）时，到期时间必须
|
|||
|
|
// 跟着取同一条记录，否则会变成「workspace 套餐 + 个人到期时间」的另一种错配。
|
|||
|
|
func TestEnrichTokenInfo_WorkspacePlanTypeKeepsItsOwnExpiry(t *testing.T) {
|
|||
|
|
const workspaceAccountID = "workspace-b"
|
|||
|
|
workspaceExpiresAt := time.Now().Add(720 * time.Hour).UTC().Format(time.RFC3339)
|
|||
|
|
|
|||
|
|
subscriptionCalls := 0
|
|||
|
|
server := newChatGPTBackendTestServer(t, chatGPTBackendTestServerConfig{
|
|||
|
|
accountsCheck: map[string]any{
|
|||
|
|
"accounts": map[string]any{
|
|||
|
|
workspaceAccountID: map[string]any{
|
|||
|
|
"account": map[string]any{
|
|||
|
|
"account_id": workspaceAccountID,
|
|||
|
|
"plan_type": "self_serve_business_usage_based",
|
|||
|
|
"is_default": true,
|
|||
|
|
},
|
|||
|
|
"entitlement": map[string]any{"expires_at": workspaceExpiresAt},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
onSubscription: func(string) map[string]any {
|
|||
|
|
subscriptionCalls++
|
|||
|
|
return map[string]any{}
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
tokenInfo := &OpenAITokenInfo{
|
|||
|
|
AccessToken: "access-token",
|
|||
|
|
ChatGPTAccountID: "personal-account-a",
|
|||
|
|
OrganizationID: workspaceAccountID,
|
|||
|
|
// id_token 没带 chatgpt_plan_type
|
|||
|
|
}
|
|||
|
|
svc := &OpenAIOAuthService{privacyClientFactory: newTestPrivacyClientFactory()}
|
|||
|
|
svc.enrichTokenInfo(context.Background(), tokenInfo, "")
|
|||
|
|
|
|||
|
|
require.Equal(t, "self_serve_business_usage_based", tokenInfo.PlanType)
|
|||
|
|
require.Equal(t, workspaceExpiresAt, tokenInfo.SubscriptionExpiresAt,
|
|||
|
|
"套餐与到期时间必须来自同一条记录")
|
|||
|
|
require.Zero(t, subscriptionCalls)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type chatGPTBackendTestServerConfig struct {
|
|||
|
|
accountsCheck map[string]any
|
|||
|
|
onSubscription func(accountID string) map[string]any
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// newChatGPTBackendTestServer 同时接管 accounts/check 与 subscriptions 两个端点，
|
|||
|
|
// 并在 t.Cleanup 里还原包级 URL 变量。
|
|||
|
|
func newChatGPTBackendTestServer(t *testing.T, cfg chatGPTBackendTestServerConfig) *httptest.Server {
|
|||
|
|
t.Helper()
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|||
|
|
w.Header().Set("Content-Type", "application/json")
|
|||
|
|
switch r.URL.Path {
|
|||
|
|
case "/backend-api/accounts/check/v4-2023-04-27":
|
|||
|
|
_ = json.NewEncoder(w).Encode(cfg.accountsCheck)
|
|||
|
|
case "/backend-api/subscriptions":
|
|||
|
|
body := map[string]any{}
|
|||
|
|
if cfg.onSubscription != nil {
|
|||
|
|
body = cfg.onSubscription(r.URL.Query().Get("account_id"))
|
|||
|
|
}
|
|||
|
|
_ = json.NewEncoder(w).Encode(body)
|
|||
|
|
default:
|
|||
|
|
w.WriteHeader(http.StatusNotFound)
|
|||
|
|
}
|
|||
|
|
}))
|
|||
|
|
|
|||
|
|
oldAccounts, oldSubscriptions := chatGPTAccountsCheckURL, chatGPTSubscriptionsURL
|
|||
|
|
chatGPTAccountsCheckURL = server.URL + "/backend-api/accounts/check/v4-2023-04-27"
|
|||
|
|
chatGPTSubscriptionsURL = server.URL + "/backend-api/subscriptions"
|
|||
|
|
t.Cleanup(func() {
|
|||
|
|
chatGPTAccountsCheckURL, chatGPTSubscriptionsURL = oldAccounts, oldSubscriptions
|
|||
|
|
})
|
|||
|
|
return server
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// enrichTokenInfo 收尾还会调用 disableOpenAITraining，它的 URL 是常量、指向真实
|
|||
|
|
// chatgpt.com，测试无法接管。给客户端一个短超时让它快速失败——该调用只写
|
|||
|
|
// PrivacyMode，不影响本组用例的断言。
|
|||
|
|
func newTestPrivacyClientFactory() PrivacyClientFactory {
|
|||
|
|
return func(string) (*req.Client, error) {
|
|||
|
|
return req.C().SetTimeout(time.Second), nil
|
|||
|
|
}
|
|||
|
|
}
|