wehub-resource-sync
493 lines
18 KiB
Swift
493 lines
18 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
@testable import CodexBar
|
|
|
|
@MainActor
|
|
final class CodexAccountPromotionTestContainer {
|
|
let suiteName: String
|
|
let rootURL: URL
|
|
let liveHomeURL: URL
|
|
let managedHomesURL: URL
|
|
let managedStoreURL: URL
|
|
let settings: SettingsStore
|
|
let usageStore: UsageStore
|
|
let fileStore: FileManagedCodexAccountStore
|
|
let homeFactory: ManagedCodexHomeFactory
|
|
let identityReader: DefaultManagedCodexIdentityReader
|
|
let workspaceResolver: any ManagedCodexWorkspaceResolving
|
|
let baseEnvironment: [String: String]
|
|
|
|
init(
|
|
suiteName: String,
|
|
workspaceIdentities: [String: CodexOpenAIWorkspaceIdentity] = [:]) throws
|
|
{
|
|
self.suiteName = suiteName
|
|
self.rootURL = FileManager.default.temporaryDirectory
|
|
.appendingPathComponent("codex-account-promotion-tests", isDirectory: true)
|
|
.appendingPathComponent(UUID().uuidString, isDirectory: true)
|
|
self.liveHomeURL = self.rootURL.appendingPathComponent("liveHome", isDirectory: true)
|
|
self.managedHomesURL = self.rootURL.appendingPathComponent("managed-codex-homes", isDirectory: true)
|
|
self.managedStoreURL = self.rootURL.appendingPathComponent("managed-codex-accounts.json", isDirectory: false)
|
|
self.baseEnvironment = ["CODEX_HOME": self.liveHomeURL.path]
|
|
self.fileStore = FileManagedCodexAccountStore(fileURL: self.managedStoreURL, fileManager: .default)
|
|
self.homeFactory = ManagedCodexHomeFactory(root: self.managedHomesURL, fileManager: .default)
|
|
self.identityReader = DefaultManagedCodexIdentityReader()
|
|
self.workspaceResolver = StubManagedCodexWorkspaceResolver(identities: workspaceIdentities)
|
|
|
|
try FileManager.default.createDirectory(at: self.liveHomeURL, withIntermediateDirectories: true)
|
|
try FileManager.default.createDirectory(at: self.managedHomesURL, withIntermediateDirectories: true)
|
|
_ = try self.fileStore.ensureFileExists()
|
|
|
|
let defaults = UserDefaults(suiteName: suiteName)!
|
|
defaults.removePersistentDomain(forName: suiteName)
|
|
defaults.set(true, forKey: "providerDetectionCompleted")
|
|
self.settings = SettingsStore(
|
|
userDefaults: defaults,
|
|
configStore: testConfigStore(suiteName: suiteName),
|
|
zaiTokenStore: NoopZaiTokenStore(),
|
|
syntheticTokenStore: NoopSyntheticTokenStore())
|
|
self.settings._test_activeManagedCodexAccount = nil
|
|
self.settings._test_activeManagedCodexRemoteHomePath = nil
|
|
self.settings._test_unreadableManagedCodexAccountStore = false
|
|
self.settings._test_managedCodexAccountStoreURL = self.managedStoreURL
|
|
self.settings._test_liveSystemCodexAccount = nil
|
|
self.settings._test_codexReconciliationEnvironment = self.baseEnvironment
|
|
self.settings.providerDetectionCompleted = true
|
|
self.settings.refreshFrequency = .manual
|
|
self.settings.codexCookieSource = .off
|
|
|
|
self.usageStore = UsageStore(
|
|
fetcher: UsageFetcher(environment: [:]),
|
|
browserDetection: BrowserDetection(cacheTTL: 0),
|
|
settings: self.settings,
|
|
startupBehavior: .testing)
|
|
self.installDynamicCodexUsageLoader()
|
|
self.usageStore._test_codexCreditsLoaderOverride = {
|
|
CreditsSnapshot(remaining: 17, events: [], updatedAt: Date())
|
|
}
|
|
}
|
|
|
|
func tearDown() {
|
|
self.usageStore._test_codexCreditsLoaderOverride = nil
|
|
self.settings._test_activeManagedCodexAccount = nil
|
|
self.settings._test_activeManagedCodexRemoteHomePath = nil
|
|
self.settings._test_unreadableManagedCodexAccountStore = false
|
|
self.settings._test_managedCodexAccountStoreURL = nil
|
|
self.settings._test_liveSystemCodexAccount = nil
|
|
self.settings._test_codexReconciliationEnvironment = nil
|
|
self.settings.userDefaults.removePersistentDomain(forName: self.suiteName)
|
|
try? FileManager.default.removeItem(at: self.rootURL)
|
|
}
|
|
|
|
func makeService(
|
|
store: (any ManagedCodexAccountStoring)? = nil,
|
|
liveAuthSwapper: (any CodexLiveAuthSwapping)? = nil,
|
|
activeSourceWriter: (any CodexActiveSourceWriting)? = nil,
|
|
snapshotLoader: (any CodexAccountReconciliationSnapshotLoading)? = nil,
|
|
accountScopedRefresher: (any CodexAccountScopedRefreshing)? = nil)
|
|
-> CodexAccountPromotionService
|
|
{
|
|
CodexAccountPromotionService(
|
|
store: store ?? self.fileStore,
|
|
homeFactory: self.homeFactory,
|
|
identityReader: self.identityReader,
|
|
workspaceResolver: self.workspaceResolver,
|
|
snapshotLoader: snapshotLoader
|
|
?? SettingsStoreCodexAccountReconciliationSnapshotLoader(settingsStore: self.settings),
|
|
authMaterialReader: DefaultCodexAuthMaterialReader(),
|
|
liveAuthSwapper: liveAuthSwapper ?? DefaultCodexLiveAuthSwapper(),
|
|
activeSourceWriter: activeSourceWriter
|
|
?? SettingsStoreCodexActiveSourceWriter(settingsStore: self.settings),
|
|
accountScopedRefresher: accountScopedRefresher
|
|
?? UsageStoreCodexAccountScopedRefresher(usageStore: self.usageStore),
|
|
baseEnvironment: self.baseEnvironment,
|
|
fileManager: .default)
|
|
}
|
|
|
|
func installDynamicCodexUsageLoader(usedPercent: Double = 12) {
|
|
let baseSpec = self.usageStore.providerSpecs[.codex]!
|
|
let liveHomePath = self.liveHomeURL.path
|
|
let identityReader = self.identityReader
|
|
self.usageStore
|
|
.providerSpecs[.codex] = makeCodexProviderSpec(baseSpec: baseSpec) {
|
|
let liveEmail = (try? identityReader.loadAccountIdentity(homePath: liveHomePath).email)
|
|
?? "unknown@example.com"
|
|
return UsageSnapshot(
|
|
primary: RateWindow(
|
|
usedPercent: usedPercent,
|
|
windowMinutes: 300,
|
|
resetsAt: nil,
|
|
resetDescription: nil),
|
|
secondary: nil,
|
|
updatedAt: Date(),
|
|
identity: ProviderIdentitySnapshot(
|
|
providerID: .codex,
|
|
accountEmail: liveEmail,
|
|
accountOrganization: nil,
|
|
loginMethod: "Pro"))
|
|
}
|
|
}
|
|
|
|
@discardableResult
|
|
func createManagedAccount(
|
|
id: UUID = UUID(),
|
|
persistedEmail: String,
|
|
authEmail: String? = nil,
|
|
authAccountID: String? = nil,
|
|
persistedProviderAccountID: String? = nil,
|
|
useAuthAccountIDAsPersistedProviderAccountID: Bool = true,
|
|
workspaceLabel: String? = nil,
|
|
workspaceAccountID: String? = nil,
|
|
plan: String = "Pro") throws -> ManagedCodexAccount
|
|
{
|
|
let homeURL = self.managedHomesURL.appendingPathComponent(id.uuidString, isDirectory: true)
|
|
let createdAt = Date().timeIntervalSince1970
|
|
let authData = try self.writeOAuthAuthFile(
|
|
homeURL: homeURL,
|
|
email: authEmail ?? persistedEmail,
|
|
plan: plan,
|
|
accountID: authAccountID)
|
|
let persistedProviderAccountIDValue: String? =
|
|
if useAuthAccountIDAsPersistedProviderAccountID {
|
|
persistedProviderAccountID ?? authAccountID
|
|
} else {
|
|
persistedProviderAccountID
|
|
}
|
|
return ManagedCodexAccount(
|
|
id: id,
|
|
email: persistedEmail,
|
|
providerAccountID: persistedProviderAccountIDValue,
|
|
workspaceLabel: workspaceLabel,
|
|
workspaceAccountID: workspaceAccountID ?? authAccountID,
|
|
authFingerprint: CodexAuthFingerprint.fingerprint(data: authData),
|
|
managedHomePath: homeURL.path,
|
|
createdAt: createdAt,
|
|
updatedAt: createdAt,
|
|
lastAuthenticatedAt: createdAt)
|
|
}
|
|
|
|
func persistAccounts(_ accounts: [ManagedCodexAccount]) throws {
|
|
try self.fileStore.storeAccounts(ManagedCodexAccountSet(
|
|
version: FileManagedCodexAccountStore.currentVersion,
|
|
accounts: accounts))
|
|
}
|
|
|
|
func loadAccounts() throws -> ManagedCodexAccountSet {
|
|
try self.fileStore.loadAccounts()
|
|
}
|
|
|
|
@discardableResult
|
|
func writeLiveOAuthAuthFile(
|
|
email: String,
|
|
plan: String = "Pro",
|
|
accountID: String? = nil,
|
|
apiKey: String? = nil) throws -> Data
|
|
{
|
|
try self.writeOAuthAuthFile(
|
|
homeURL: self.liveHomeURL,
|
|
email: email,
|
|
plan: plan,
|
|
accountID: accountID,
|
|
apiKey: apiKey)
|
|
}
|
|
|
|
@discardableResult
|
|
func writeLiveAPIKeyAuthFile(apiKey: String = "sk-live-only") throws -> Data {
|
|
let data = try JSONSerialization.data(
|
|
withJSONObject: ["OPENAI_API_KEY": apiKey],
|
|
options: [.sortedKeys])
|
|
try FileManager.default.createDirectory(at: self.liveHomeURL, withIntermediateDirectories: true)
|
|
try data.write(to: Self.authFileURL(for: self.liveHomeURL), options: .atomic)
|
|
return data
|
|
}
|
|
|
|
@discardableResult
|
|
func writeLiveOAuthAuthFileWithoutEmail(accountID: String, apiKey: String? = nil) throws -> Data {
|
|
try self.writeOAuthAuthFileWithoutEmail(
|
|
homeURL: self.liveHomeURL,
|
|
accountID: accountID,
|
|
apiKey: apiKey)
|
|
}
|
|
|
|
@discardableResult
|
|
func writeManagedOAuthAuthFileWithoutEmail(
|
|
for account: ManagedCodexAccount,
|
|
accountID: String,
|
|
apiKey: String? = nil) throws -> Data
|
|
{
|
|
try self.writeOAuthAuthFileWithoutEmail(
|
|
homeURL: URL(fileURLWithPath: account.managedHomePath, isDirectory: true),
|
|
accountID: accountID,
|
|
apiKey: apiKey)
|
|
}
|
|
|
|
@discardableResult
|
|
private func writeOAuthAuthFileWithoutEmail(
|
|
homeURL: URL,
|
|
accountID: String,
|
|
apiKey: String? = nil) throws -> Data
|
|
{
|
|
try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true)
|
|
|
|
let tokens: [String: Any] = [
|
|
"accessToken": "access-\(accountID)",
|
|
"refreshToken": "refresh-\(accountID)",
|
|
"accountId": accountID,
|
|
]
|
|
var json: [String: Any] = [
|
|
"tokens": tokens,
|
|
"last_refresh": "2026-04-05T00:00:00Z",
|
|
]
|
|
if let apiKey {
|
|
json["OPENAI_API_KEY"] = apiKey
|
|
}
|
|
|
|
let data = try JSONSerialization.data(withJSONObject: json, options: [.sortedKeys])
|
|
try data.write(to: Self.authFileURL(for: homeURL), options: .atomic)
|
|
return data
|
|
}
|
|
|
|
func removeLiveAuthFile() throws {
|
|
let authFileURL = Self.authFileURL(for: self.liveHomeURL)
|
|
if FileManager.default.fileExists(atPath: authFileURL.path) {
|
|
try FileManager.default.removeItem(at: authFileURL)
|
|
}
|
|
}
|
|
|
|
func liveAuthData() throws -> Data? {
|
|
let authFileURL = Self.authFileURL(for: self.liveHomeURL)
|
|
guard FileManager.default.fileExists(atPath: authFileURL.path) else { return nil }
|
|
return try Data(contentsOf: authFileURL)
|
|
}
|
|
|
|
func managedAuthData(for account: ManagedCodexAccount) throws -> Data {
|
|
try Data(contentsOf: Self.authFileURL(for: URL(fileURLWithPath: account.managedHomePath, isDirectory: true)))
|
|
}
|
|
|
|
func managedHomeURLs() throws -> [URL] {
|
|
let urls = try FileManager.default.contentsOfDirectory(
|
|
at: self.managedHomesURL,
|
|
includingPropertiesForKeys: nil)
|
|
return urls.sorted { $0.lastPathComponent < $1.lastPathComponent }
|
|
}
|
|
|
|
func seedScopedRefreshState(email: String, identity: CodexIdentity) {
|
|
let snapshot = UsageSnapshot(
|
|
primary: RateWindow(usedPercent: 4, windowMinutes: 300, resetsAt: nil, resetDescription: nil),
|
|
secondary: nil,
|
|
updatedAt: Date(),
|
|
identity: ProviderIdentitySnapshot(
|
|
providerID: .codex,
|
|
accountEmail: email,
|
|
accountOrganization: nil,
|
|
loginMethod: "Pro"))
|
|
let credits = CreditsSnapshot(remaining: 3, events: [], updatedAt: Date())
|
|
|
|
self.usageStore._setSnapshotForTesting(snapshot, provider: .codex)
|
|
self.usageStore.credits = credits
|
|
self.usageStore.lastCreditsSnapshot = credits
|
|
self.usageStore.lastCreditsSnapshotAccountKey = email
|
|
self.usageStore.lastCreditsSource = .api
|
|
self.usageStore.lastCodexAccountScopedRefreshGuard = CodexAccountScopedRefreshGuard(
|
|
source: .liveSystem,
|
|
identity: identity,
|
|
accountKey: email)
|
|
}
|
|
|
|
@discardableResult
|
|
private func writeOAuthAuthFile(
|
|
homeURL: URL,
|
|
email: String,
|
|
plan: String,
|
|
accountID: String?,
|
|
apiKey: String? = nil) throws -> Data
|
|
{
|
|
try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true)
|
|
|
|
var tokens: [String: Any] = [
|
|
"accessToken": "access-\(email)",
|
|
"refreshToken": "refresh-\(email)",
|
|
"idToken": Self.fakeJWT(email: email, plan: plan, accountID: accountID),
|
|
]
|
|
if let accountID {
|
|
tokens["accountId"] = accountID
|
|
}
|
|
|
|
var json: [String: Any] = [
|
|
"tokens": tokens,
|
|
"last_refresh": "2026-04-05T00:00:00Z",
|
|
]
|
|
if let apiKey {
|
|
json["OPENAI_API_KEY"] = apiKey
|
|
}
|
|
|
|
let data = try JSONSerialization.data(withJSONObject: json, options: [.sortedKeys])
|
|
try data.write(to: Self.authFileURL(for: homeURL), options: .atomic)
|
|
return data
|
|
}
|
|
|
|
private static func fakeJWT(email: String, plan: String, accountID: String?) -> String {
|
|
let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data()
|
|
var authClaims: [String: Any] = [
|
|
"chatgpt_plan_type": plan,
|
|
]
|
|
if let accountID {
|
|
authClaims["chatgpt_account_id"] = accountID
|
|
}
|
|
let payload = (try? JSONSerialization.data(withJSONObject: [
|
|
"email": email,
|
|
"chatgpt_plan_type": plan,
|
|
"https://api.openai.com/auth": authClaims,
|
|
])) ?? Data()
|
|
|
|
func base64URL(_ data: Data) -> String {
|
|
data.base64EncodedString()
|
|
.replacingOccurrences(of: "=", with: "")
|
|
.replacingOccurrences(of: "+", with: "-")
|
|
.replacingOccurrences(of: "/", with: "_")
|
|
}
|
|
|
|
return "\(base64URL(header)).\(base64URL(payload))."
|
|
}
|
|
|
|
private static func authFileURL(for homeURL: URL) -> URL {
|
|
homeURL.appendingPathComponent("auth.json", isDirectory: false)
|
|
}
|
|
}
|
|
|
|
private func makeCodexProviderSpec(
|
|
baseSpec: ProviderSpec,
|
|
loader: @escaping @Sendable () async throws -> UsageSnapshot) -> ProviderSpec
|
|
{
|
|
let baseDescriptor = baseSpec.descriptor
|
|
let strategy = TestPromotionCodexFetchStrategy(loader: loader)
|
|
let descriptor = ProviderDescriptor(
|
|
id: .codex,
|
|
metadata: baseDescriptor.metadata,
|
|
branding: baseDescriptor.branding,
|
|
tokenCost: baseDescriptor.tokenCost,
|
|
fetchPlan: ProviderFetchPlan(
|
|
sourceModes: [.auto, .cli, .oauth],
|
|
pipeline: ProviderFetchPipeline { _ in [strategy] }),
|
|
cli: baseDescriptor.cli)
|
|
return ProviderSpec(
|
|
style: baseSpec.style,
|
|
isEnabled: baseSpec.isEnabled,
|
|
descriptor: descriptor,
|
|
makeFetchContext: baseSpec.makeFetchContext)
|
|
}
|
|
|
|
private struct TestPromotionCodexFetchStrategy: ProviderFetchStrategy {
|
|
let loader: @Sendable () async throws -> UsageSnapshot
|
|
|
|
var id: String {
|
|
"test-promotion-codex"
|
|
}
|
|
|
|
var kind: ProviderFetchKind {
|
|
.cli
|
|
}
|
|
|
|
func isAvailable(_: ProviderFetchContext) async -> Bool {
|
|
true
|
|
}
|
|
|
|
func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult {
|
|
let snapshot = try await self.loader()
|
|
return self.makeResult(usage: snapshot, sourceLabel: "test-promotion-codex")
|
|
}
|
|
|
|
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
|
|
false
|
|
}
|
|
}
|
|
|
|
final class RecordingManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable {
|
|
let base: any ManagedCodexAccountStoring
|
|
var storedSnapshots: [ManagedCodexAccountSet] = []
|
|
var onStore: (@Sendable (ManagedCodexAccountSet) throws -> Void)?
|
|
|
|
init(
|
|
base: any ManagedCodexAccountStoring,
|
|
onStore: (@Sendable (ManagedCodexAccountSet) throws -> Void)? = nil)
|
|
{
|
|
self.base = base
|
|
self.onStore = onStore
|
|
}
|
|
|
|
func loadAccounts() throws -> ManagedCodexAccountSet {
|
|
try self.base.loadAccounts()
|
|
}
|
|
|
|
func storeAccounts(_ accounts: ManagedCodexAccountSet) throws {
|
|
self.storedSnapshots.append(accounts)
|
|
try self.onStore?(accounts)
|
|
try self.base.storeAccounts(accounts)
|
|
}
|
|
|
|
func ensureFileExists() throws -> URL {
|
|
try self.base.ensureFileExists()
|
|
}
|
|
}
|
|
|
|
final class RecordingCodexLiveAuthSwapper: CodexLiveAuthSwapping, @unchecked Sendable {
|
|
let base: any CodexLiveAuthSwapping
|
|
var swapCallCount = 0
|
|
var swappedData: [Data] = []
|
|
var onSwap: (@Sendable (Data, URL) throws -> Void)?
|
|
|
|
init(
|
|
base: any CodexLiveAuthSwapping = DefaultCodexLiveAuthSwapper(),
|
|
onSwap: (@Sendable (Data, URL) throws -> Void)? = nil)
|
|
{
|
|
self.base = base
|
|
self.onSwap = onSwap
|
|
}
|
|
|
|
func swapLiveAuthData(_ data: Data, liveHomeURL: URL) throws {
|
|
self.swapCallCount += 1
|
|
self.swappedData.append(data)
|
|
try self.onSwap?(data, liveHomeURL)
|
|
try self.base.swapLiveAuthData(data, liveHomeURL: liveHomeURL)
|
|
}
|
|
}
|
|
|
|
enum PromotionTestError: Error, Equatable {
|
|
case storeWriteFailed
|
|
case swapFailed
|
|
case unexpectedDisposition
|
|
}
|
|
|
|
private struct StubManagedCodexWorkspaceResolver: ManagedCodexWorkspaceResolving {
|
|
let identities: [String: CodexOpenAIWorkspaceIdentity]
|
|
|
|
func resolveWorkspaceIdentity(
|
|
homePath _: String,
|
|
providerAccountID: String) async -> CodexOpenAIWorkspaceIdentity?
|
|
{
|
|
self.identities[providerAccountID]
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
struct StaticCodexAccountReconciliationSnapshotLoader: CodexAccountReconciliationSnapshotLoading {
|
|
let snapshot: CodexAccountReconciliationSnapshot
|
|
|
|
func loadSnapshot() -> CodexAccountReconciliationSnapshot {
|
|
self.snapshot
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
final class ClosureCodexAccountScopedRefresher: CodexAccountScopedRefreshing {
|
|
private let onRefresh: @MainActor (Bool) async -> Void
|
|
|
|
init(onRefresh: @escaping @MainActor (Bool) async -> Void) {
|
|
self.onRefresh = onRefresh
|
|
}
|
|
|
|
func refreshCodexAccountScopedState(allowDisabled: Bool) async {
|
|
await self.onRefresh(allowDisabled)
|
|
}
|
|
}
|