Peter Steinbergertest: stabilize release prechecka6a538e5 months agoHistory
364 lines · 14.9 KB · Swift
RawDownload
import Foundationimport Testing@testable import CodexBar@testable import CodexBarCore @Suite(.serialized)@MainActorstruct CodexBackgroundRefreshCoalescingTests {    @Test    func `rapid regular refreshes coalesce concurrent Codex credits fetches`() async throws {        let settings = try self.makeSettingsStore(            suite: "CodexBackgroundRefreshCoalescingTests-credits-coalescing")        settings.statusChecksEnabled = false        settings.openAIWebAccessEnabled = false        let managedAccount = try Self.installManagedAccount(            email: "managed@example.com",            settings: settings)        defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) }         let store = self.makeStore(settings: settings)        let blocker = BlockingCreditsLoader()        let firstCompletion = RefreshCompletionProbe()        let secondCompletion = RefreshCompletionProbe()        store._test_providerRefreshOverride = { _ in }        defer { store._test_providerRefreshOverride = nil }        store._test_codexCreditsLoaderOverride = {            try await blocker.awaitResult()        }        defer { store._test_codexCreditsLoaderOverride = nil }         let firstRefreshTask = Task {            await store.refresh(forceTokenUsage: false)            await firstCompletion.markCompleted()        }        await blocker.waitUntilStarted(count: 1)        #expect(await firstCompletion.waitUntilCompleted() == true)         let secondRefreshTask = Task {            await store.refresh(forceTokenUsage: false)            await secondCompletion.markCompleted()        }         #expect(await secondCompletion.waitUntilCompleted() == true)        #expect(await blocker.startedCount() == 1)         await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date())))         await firstRefreshTask.value        await secondRefreshTask.value    }     @Test    func `regular credits refresh reschedules when Codex account changes`() async throws {        let settings = try self.makeSettingsStore(            suite: "CodexBackgroundRefreshCoalescingTests-credits-account-switch")        settings.statusChecksEnabled = false        settings.openAIWebAccessEnabled = false        let alphaAccount = try Self.makeManagedAccount(email: "alpha@example.com")        let betaAccount = try Self.makeManagedAccount(email: "beta@example.com")        defer {            try? FileManager.default.removeItem(atPath: alphaAccount.managedHomePath)            try? FileManager.default.removeItem(atPath: betaAccount.managedHomePath)        }        settings._test_activeManagedCodexAccount = alphaAccount        settings.codexActiveSource = .managedAccount(id: alphaAccount.id)        defer { settings._test_activeManagedCodexAccount = nil }         let store = self.makeStore(settings: settings)        let blocker = BlockingCreditsLoader()        store._test_providerRefreshOverride = { _ in }        defer { store._test_providerRefreshOverride = nil }        store._test_codexCreditsLoaderOverride = {            try await blocker.awaitResult()        }        defer { store._test_codexCreditsLoaderOverride = nil }         let alphaRefreshTask = Task {            await store.refresh(forceTokenUsage: false)        }        await blocker.waitUntilStarted(count: 1)         settings._test_activeManagedCodexAccount = betaAccount        settings.codexActiveSource = .managedAccount(id: betaAccount.id)        let betaRefreshTask = Task {            await store.refresh(forceTokenUsage: false)        }        await blocker.waitUntilStarted(count: 2)         await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 10, events: [], updatedAt: Date())))        await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date())))         await alphaRefreshTask.value        await betaRefreshTask.value        await store.creditsRefreshTask?.value         #expect(await blocker.startedCount() == 2)        #expect(store.lastCreditsSnapshotAccountKey == "beta@example.com")        #expect(store.credits?.remaining == 25)    }     @Test    func `force refresh cancels stale background Codex credits fetch`() async throws {        let settings = try self.makeSettingsStore(            suite: "CodexBackgroundRefreshCoalescingTests-credits-force-cancels-background")        settings.statusChecksEnabled = false        settings.openAIWebAccessEnabled = false        let managedAccount = try Self.installManagedAccount(            email: "managed@example.com",            settings: settings)        defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) }         let store = self.makeStore(settings: settings)        let blocker = BlockingCreditsLoader()        store._test_providerRefreshOverride = { _ in }        defer { store._test_providerRefreshOverride = nil }        store._test_codexCreditsLoaderOverride = {            try await blocker.awaitResult()        }        defer { store._test_codexCreditsLoaderOverride = nil }        let regularCompletion = RefreshCompletionProbe()         let regularRefreshTask = Task {            await store.refresh(forceTokenUsage: false)            await regularCompletion.markCompleted()        }        await blocker.waitUntilStarted(count: 1)        #expect(await regularCompletion.waitUntilCompleted() == true)         let forceRefreshTask = Task {            await store.refresh(forceTokenUsage: true)        }        await blocker.waitUntilStarted(count: 2)         await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 10, events: [], updatedAt: Date())))        await blocker.resumeNext(with: .success(CreditsSnapshot(remaining: 25, events: [], updatedAt: Date())))         await regularRefreshTask.value        await forceRefreshTask.value         #expect(await blocker.startedCount() == 2)        #expect(store.credits?.remaining == 25)    }     @Test    func `rapid regular refreshes coalesce concurrent OpenAI dashboard fetches`() async throws {        let settings = try self.makeSettingsStore(            suite: "CodexBackgroundRefreshCoalescingTests-dashboard-coalescing")        settings.statusChecksEnabled = false        let managedAccount = try Self.installManagedAccount(            email: "managed@example.com",            settings: settings)        defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) }         let store = self.makeStore(settings: settings)        let blocker = BlockingManagedOpenAIDashboardLoader()        let firstCompletion = RefreshCompletionProbe()        let secondCompletion = RefreshCompletionProbe()        store._test_providerRefreshOverride = { _ in }        defer { store._test_providerRefreshOverride = nil }        store._test_codexCreditsLoaderOverride = {            CreditsSnapshot(remaining: 25, events: [], updatedAt: Date())        }        defer { store._test_codexCreditsLoaderOverride = nil }        await store.refresh(forceTokenUsage: false)         store._test_openAIDashboardLoaderOverride = { _, _, _, _ in            try await blocker.awaitResult()        }        defer { store._test_openAIDashboardLoaderOverride = nil }         let firstRefreshTask = Task {            await store.refresh(forceTokenUsage: false)            await firstCompletion.markCompleted()        }        await blocker.waitUntilStarted(count: 1)        #expect(await firstCompletion.waitUntilCompleted() == true)         let secondRefreshTask = Task {            await store.refresh(forceTokenUsage: false)            await secondCompletion.markCompleted()        }         #expect(await secondCompletion.waitUntilCompleted() == true)        #expect(await blocker.startedCount() == 1)         let backgroundTask = try #require(store.openAIDashboardBackgroundRefreshTask)        await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot(            signedInEmail: managedAccount.email,            codeReviewRemainingPercent: 95,            creditEvents: [],            dailyBreakdown: [],            usageBreakdown: [],            creditsPurchaseURL: nil,            creditsRemaining: 25,            accountPlan: "Pro",            updatedAt: Date())))         await firstRefreshTask.value        await secondRefreshTask.value        await backgroundTask.value         #expect(store.openAIDashboard?.creditsRemaining == 25)    }     @Test    func `cancelled background dashboard import does not publish stale account status`() async throws {        let settings = try self.makeSettingsStore(            suite: "CodexBackgroundRefreshCoalescingTests-dashboard-cancelled-import")        settings.statusChecksEnabled = false        let managedAccount = try Self.installManagedAccount(            email: "managed@example.com",            settings: settings)        defer { try? FileManager.default.removeItem(atPath: managedAccount.managedHomePath) }         let store = self.makeStore(settings: settings)        let importBlocker = BlockingOpenAIDashboardCookieImport()        store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in            try await importBlocker.awaitResult()        }        defer { store._test_openAIDashboardCookieImportOverride = nil }         let importTask = Task { @MainActor in            await store.importOpenAIDashboardCookiesIfNeeded(                targetEmail: managedAccount.email,                force: true)        }        await importBlocker.waitUntilStarted()        importTask.cancel()        await importBlocker.resumeNext(with: .failure(            OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount(                found: [.init(sourceLabel: "Chrome", email: "other@example.com")])))         let imported = await importTask.value        #expect(imported == nil)        #expect(store.openAIDashboard == nil)        #expect(store.openAIDashboardCookieImportStatus == nil)        #expect(store.openAIDashboardRequiresLogin == false)    }     private func makeSettingsStore(suite: String) throws -> SettingsStore {        let defaults = UserDefaults(suiteName: suite)!        defaults.removePersistentDomain(forName: suite)        defaults.set(true, forKey: "providerDetectionCompleted")        let configStore = testConfigStore(suiteName: suite)        let settings = SettingsStore(            userDefaults: defaults,            configStore: configStore,            zaiTokenStore: NoopZaiTokenStore(),            syntheticTokenStore: NoopSyntheticTokenStore())        let codexMetadata = try #require(ProviderDescriptorRegistry.metadata[.codex])        settings.setProviderEnabled(provider: .codex, metadata: codexMetadata, enabled: true)        settings.providerDetectionCompleted = true        settings.openAIWebAccessEnabled = true        settings.codexCookieSource = .auto        return settings    }     private func makeStore(settings: SettingsStore) -> UsageStore {        UsageStore(            fetcher: UsageFetcher(environment: [:]),            browserDetection: BrowserDetection(cacheTTL: 0),            settings: settings,            startupBehavior: .testing)    }     private static func installManagedAccount(        email: String,        settings: SettingsStore) throws -> ManagedCodexAccount    {        let account = try Self.makeManagedAccount(email: email)        settings._test_activeManagedCodexAccount = account        settings.codexActiveSource = .managedAccount(id: account.id)        return account    }     private static func makeManagedAccount(email: String) throws -> ManagedCodexAccount {        let managedHomeURL = FileManager.default.temporaryDirectory            .appendingPathComponent(UUID().uuidString, isDirectory: true)        try Self.writeCodexAuthFile(            homeURL: managedHomeURL,            email: email,            plan: "Pro")        return ManagedCodexAccount(            id: UUID(),            email: email,            managedHomePath: managedHomeURL.path,            createdAt: 1,            updatedAt: 1,            lastAuthenticatedAt: 1)    }     private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws {        try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true)        let tokens: [String: Any] = [            "accessToken": "access-token",            "refreshToken": "refresh-token",            "idToken": Self.fakeJWT(email: email, plan: plan),        ]        let data = try JSONSerialization.data(withJSONObject: ["tokens": tokens], options: [.sortedKeys])        try data.write(to: homeURL.appendingPathComponent("auth.json"))    }     private static func fakeJWT(email: String, plan: String) -> String {        let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data()        let payload = (try? JSONSerialization.data(withJSONObject: [            "email": email,            "chatgpt_plan_type": plan,            "https://api.openai.com/auth": [                "chatgpt_plan_type": plan,            ],        ])) ?? Data()         func base64URL(_ data: Data) -> String {            data.base64EncodedString()                .replacingOccurrences(of: "=", with: "")                .replacingOccurrences(of: "+", with: "-")                .replacingOccurrences(of: "/", with: "_")        }         return "\(base64URL(header)).\(base64URL(payload))."    }} private actor BlockingOpenAIDashboardCookieImport {    private var continuations: [        CheckedContinuation<Result<OpenAIDashboardBrowserCookieImporter.ImportResult, Error>, Never>    ] = []    private var startWaiters: [(count: Int, continuation: CheckedContinuation<Void, Never>)] = []    private var started = 0     func awaitResult() async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult {        let result = await withCheckedContinuation { continuation in            self.continuations.append(continuation)            self.started += 1            self.resumeReadyStartWaiters()        }        return try result.get()    }     func waitUntilStarted(count: Int = 1) async {        if self.started >= count { return }        await withCheckedContinuation { continuation in            self.startWaiters.append((count: count, continuation: continuation))        }    }     func resumeNext(with result: Result<OpenAIDashboardBrowserCookieImporter.ImportResult, Error>) {        guard !self.continuations.isEmpty else { return }        let continuation = self.continuations.removeFirst()        continuation.resume(returning: result)    }     private func resumeReadyStartWaiters() {        var remaining: [(count: Int, continuation: CheckedContinuation<Void, Never>)] = []        for waiter in self.startWaiters {            if self.started >= waiter.count {                waiter.continuation.resume()            } else {                remaining.append(waiter)            }        }        self.startWaiters = remaining    }}