346 lines · 14.0 KB · Swift
RawDownload
import AppKitimport CodexBarCoreimport Foundationimport Testing@testable import CodexBar @Suite(.serialized)@MainActorstruct CodexAccountMenuDisplaySnapshotTests {    private func makeSettings() -> SettingsStore {        let suite = "CodexAccountMenuDisplaySnapshotTests-\(UUID().uuidString)"        let defaults = UserDefaults(suiteName: suite)!        defaults.removePersistentDomain(forName: suite)        defaults.set(true, forKey: "providerDetectionCompleted")        let settings = SettingsStore(            userDefaults: defaults,            configStore: testConfigStore(suiteName: suite),            zaiTokenStore: NoopZaiTokenStore(),            syntheticTokenStore: NoopSyntheticTokenStore())        settings.providerDetectionCompleted = true        return settings    }     private func enableOnlyCodex(_ settings: SettingsStore) {        for provider in UsageProvider.allCases {            guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue }            settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex)        }    }     private func liveSnapshot(email: String) -> CodexAccountReconciliationSnapshot {        CodexAccountReconciliationSnapshot(            storedAccounts: [],            activeStoredAccount: nil,            liveSystemAccount: ObservedSystemCodexAccount(                email: email,                codexHomePath: "/tmp/\(email)",                observedAt: Date()),            matchingStoredAccountForLiveSystemAccount: nil,            activeSource: .liveSystem,            hasUnreadableAddedAccountStore: false)    }     private func cachedProjection(        snapshot: CodexAccountReconciliationSnapshot,        loadedAt: Date = Date(timeIntervalSinceNow: -3600)) -> CachedCodexAccountMenuProjection    {        CachedCodexAccountMenuProjection(            activeSource: snapshot.activeSource,            loadedAt: loadedAt,            projection: CodexVisibleAccountProjection.make(from: snapshot))    }     @Test    func `cold menu projection read never loads auth state`() async {        let settings = self.makeSettings()        let probe = CodexAccountSnapshotLoaderProbe(snapshot: self.liveSnapshot(email: "loaded@example.com"))        settings._test_codexAccountSnapshotLoader = { _ in probe.load() }        defer { settings._test_codexAccountSnapshotLoader = nil }         #expect(settings.codexVisibleAccountProjectionForMenuDisplay == nil)        #expect(probe.callCount == 0)         let result = await settings.revalidateCodexAccountMenuProjection()         #expect(result == .updated)        #expect(probe.callCount == 1)        #expect(probe.loadedOffMainThread)        #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email ==                "loaded@example.com")    }     @Test    func `override snapshot load preserves persisted account menu projection`() {        let settings = self.makeSettings()        let activeSnapshot = self.liveSnapshot(email: "active@example.com")        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: activeSnapshot)         let otherID = UUID()        let otherAccount = ManagedCodexAccount(            id: otherID,            email: "other@example.com",            managedHomePath: "/tmp/other",            createdAt: 1,            updatedAt: 1,            lastAuthenticatedAt: 1)        let overrideSnapshot = CodexAccountReconciliationSnapshot(            storedAccounts: [otherAccount],            activeStoredAccount: otherAccount,            liveSystemAccount: nil,            matchingStoredAccountForLiveSystemAccount: nil,            activeSource: .managedAccount(id: otherID),            hasUnreadableAddedAccountStore: false)        settings._test_codexAccountSnapshotLoader = { _ in overrideSnapshot }        defer { settings._test_codexAccountSnapshotLoader = nil }         _ = settings.codexAccountReconciliationSnapshot(activeSourceOverride: .managedAccount(id: otherID))         #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email ==                "active@example.com")    }     @Test    func `managed account change refreshes account menu projection`() {        let settings = self.makeSettings()        let activeSnapshot = self.liveSnapshot(email: "active@example.com")        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: activeSnapshot, loadedAt: Date())         let addedAccount = ManagedCodexAccount(            id: UUID(),            email: "added@example.com",            managedHomePath: "/tmp/added",            createdAt: 1,            updatedAt: 1,            lastAuthenticatedAt: 1)        let refreshedSnapshot = CodexAccountReconciliationSnapshot(            storedAccounts: [addedAccount],            activeStoredAccount: nil,            liveSystemAccount: activeSnapshot.liveSystemAccount,            matchingStoredAccountForLiveSystemAccount: nil,            activeSource: .liveSystem,            hasUnreadableAddedAccountStore: false)        settings._test_codexAccountSnapshotLoader = { _ in refreshedSnapshot }        defer { settings._test_codexAccountSnapshotLoader = nil }         settings.refreshCodexAccountReconciliationAfterManagedAccountsDidChange()         #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.contains {                $0.email == "added@example.com"            } == true)    }     @Test    func `stale menu projection returns immediately then refreshes concurrently`() async {        let settings = self.makeSettings()        let staleSnapshot = self.liveSnapshot(email: "before@example.com")        let probe = CodexAccountSnapshotLoaderProbe(snapshot: self.liveSnapshot(email: "after@example.com"))        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot)        settings._test_codexAccountSnapshotLoader = { _ in probe.load() }        SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60        defer {            SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil            settings._test_codexAccountSnapshotLoader = nil        }         #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email ==                "before@example.com")        #expect(probe.callCount == 0)        #expect(settings.codexAccountMenuProjectionNeedsRevalidation)         let result = await settings.revalidateCodexAccountMenuProjection()         #expect(result == .updated)        #expect(probe.callCount == 1)        #expect(probe.loadedOffMainThread)        #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email ==                "after@example.com")    }     @Test    func `revalidation discards result after reconciliation generation changes`() async {        let settings = self.makeSettings()        let staleSnapshot = self.liveSnapshot(email: "before@example.com")        let probe = CodexAccountSnapshotLoaderProbe(            snapshot: self.liveSnapshot(email: "discarded@example.com"),            blocks: true)        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot)        settings._test_codexAccountSnapshotLoader = { _ in probe.load() }        SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60        defer {            probe.release()            SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil            settings._test_codexAccountSnapshotLoader = nil        }         let task = Task { await settings.revalidateCodexAccountMenuProjection() }        await probe.waitUntilCalled()        settings.invalidateCodexAccountReconciliationSnapshotCache()        probe.release()         #expect(await task.value == .discarded)        #expect(            settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email ==                "before@example.com")    }     @Test    func `fresh menu open coalesces account projection revalidation and identity stays read only`() async throws {        StatusItemController.menuCardRenderingEnabled = false        StatusItemController.setMenuRefreshEnabledForTesting(false)        StatusItemController.setCodexAccountMenuProjectionRevalidationEnabledForTesting(true)        defer {            StatusItemController.resetCodexAccountMenuProjectionRevalidationEnabledForTesting()            StatusItemController.resetMenuRefreshEnabledForTesting()        }         let settings = self.makeSettings()        settings.statusChecksEnabled = false        settings.refreshFrequency = .manual        settings.mergeIcons = false        self.enableOnlyCodex(settings)        let store = UsageStore(            fetcher: UsageFetcher(),            browserDetection: BrowserDetection(cacheTTL: 0),            settings: settings)        let controller = StatusItemController(            store: store,            settings: settings,            account: AccountInfo(email: nil, plan: nil),            updater: DisabledUpdaterController(),            preferencesSelection: PreferencesSelection(),            statusBar: .system)        defer { controller.releaseStatusItemsForTesting() }         let staleSnapshot = self.liveSnapshot(email: "before@example.com")        let probe = CodexAccountSnapshotLoaderProbe(            snapshot: self.liveSnapshot(email: "after@example.com"),            blocks: true)        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot)        settings._test_codexAccountSnapshotLoader = { _ in probe.load() }        SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60        defer {            probe.release()            SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil            settings._test_codexAccountSnapshotLoader = nil        }         let menu = NSMenu()        controller.menuProviders[ObjectIdentifier(menu)] = .codex        controller.markMenuFresh(menu)        #expect(controller.codexAccountMenuDisplay(for: .codex) == nil)        #expect(probe.callCount == 0)         let versionBeforeOpen = controller.menuContentVersion        controller.menuWillOpen(menu)        let revalidation = try #require(controller.codexAccountMenuProjectionRevalidationTask)        controller.menuWillOpen(menu)        await probe.waitUntilCalled()         #expect(probe.callCount == 1)        #expect(probe.loadedOffMainThread)        probe.release()        await revalidation.value         #expect(controller.codexAccountMenuProjectionRevalidationTask == nil)        #expect(controller.menuContentVersion == versionBeforeOpen + 1)    }     @Test    func `selecting displayed account uses captured source without reconciliation`() throws {        let settings = self.makeSettings()        let firstID = UUID()        let secondID = UUID()        let first = ManagedCodexAccount(            id: firstID,            email: "first@example.com",            managedHomePath: "/tmp/first",            createdAt: 1,            updatedAt: 1,            lastAuthenticatedAt: 1)        let second = ManagedCodexAccount(            id: secondID,            email: "second@example.com",            managedHomePath: "/tmp/second",            createdAt: 1,            updatedAt: 1,            lastAuthenticatedAt: 1)        let snapshot = CodexAccountReconciliationSnapshot(            storedAccounts: [first, second],            activeStoredAccount: first,            liveSystemAccount: nil,            matchingStoredAccountForLiveSystemAccount: nil,            activeSource: .managedAccount(id: firstID),            hasUnreadableAddedAccountStore: false)        let projection = CodexVisibleAccountProjection.make(from: snapshot)        let displayedAccount = try #require(projection.visibleAccounts.first {            $0.selectionSource == .managedAccount(id: secondID)        })        let probe = CodexAccountSnapshotLoaderProbe(snapshot: snapshot)        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: snapshot)        settings._test_codexAccountSnapshotLoader = { _ in probe.load() }        defer { settings._test_codexAccountSnapshotLoader = nil }         settings.selectDisplayedCodexVisibleAccount(displayedAccount)         #expect(probe.callCount == 0)        #expect(settings.codexActiveSource == .managedAccount(id: secondID))        #expect(settings.cachedCodexAccountMenuProjection == nil)    }} private final class CodexAccountSnapshotLoaderProbe: @unchecked Sendable {    private let lock = NSLock()    private let snapshot: CodexAccountReconciliationSnapshot    private let blocks: Bool    private let releaseSemaphore = DispatchSemaphore(value: 0)    private var _callCount = 0    private var _loadedOffMainThread = false    private var released = false     init(snapshot: CodexAccountReconciliationSnapshot, blocks: Bool = false) {        self.snapshot = snapshot        self.blocks = blocks    }     var callCount: Int {        self.lock.withLock { self._callCount }    }     var loadedOffMainThread: Bool {        self.lock.withLock { self._loadedOffMainThread }    }     func load() -> CodexAccountReconciliationSnapshot {        self.lock.withLock {            self._callCount += 1            self._loadedOffMainThread = self._loadedOffMainThread || !Thread.isMainThread        }        if self.blocks {            self.releaseSemaphore.wait()        }        return self.snapshot    }     func waitUntilCalled() async {        while self.callCount == 0 {            await Task.yield()        }    }     func release() {        let shouldSignal = self.lock.withLock {            guard !self.released else { return false }            self.released = true            return true        }        if shouldSignal {            self.releaseSemaphore.signal()        }    }}