steipete--codexbar/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift

import AppKit
import CodexBarCore
import Foundation
import Testing
@testable import CodexBar

@Suite(.serialized)
@MainActor
struct 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 `stacked menu matches runtime enriched snapshot for legacy managed workspace`() throws {
        let settings = self.makeSettings()
        settings.multiAccountMenuLayout = .stacked
        let legacyID = UUID()
        let siblingID = UUID()
        let legacy = ManagedCodexAccount(
            id: legacyID,
            email: "legacy@example.com",
            workspaceAccountID: nil,
            managedHomePath: "/tmp/legacy",
            createdAt: 1,
            updatedAt: 1,
            lastAuthenticatedAt: 1)
        let sibling = ManagedCodexAccount(
            id: siblingID,
            email: "sibling@example.com",
            workspaceAccountID: "account-sibling",
            managedHomePath: "/tmp/sibling",
            createdAt: 1,
            updatedAt: 1,
            lastAuthenticatedAt: 1)
        let snapshot = CodexAccountReconciliationSnapshot(
            storedAccounts: [legacy, sibling],
            activeStoredAccount: legacy,
            liveSystemAccount: nil,
            matchingStoredAccountForLiveSystemAccount: nil,
            activeSource: .managedAccount(id: legacyID),
            hasUnreadableAddedAccountStore: false,
            storedAccountRuntimeIdentities: [
                legacyID: .providerAccount(id: " Account-Runtime "),
                siblingID: .providerAccount(id: "account-sibling"),
            ])
        let projection = CodexVisibleAccountProjection.make(from: snapshot)
        let legacyProjected = try #require(projection.visibleAccounts.first {
            $0.selectionSource == .managedAccount(id: legacyID)
        })
        let siblingProjected = try #require(projection.visibleAccounts.first {
            $0.selectionSource == .managedAccount(id: siblingID)
        })
        let runtimeEnrichedLegacy = CodexVisibleAccount(
            id: legacyProjected.id,
            email: legacyProjected.email,
            workspaceLabel: legacyProjected.workspaceLabel,
            workspaceAccountID: "account-runtime",
            authFingerprint: legacyProjected.authFingerprint,
            storedAccountID: legacyProjected.storedAccountID,
            selectionSource: legacyProjected.selectionSource,
            isActive: legacyProjected.isActive,
            isLive: legacyProjected.isLive,
            canReauthenticate: legacyProjected.canReauthenticate,
            canRemove: legacyProjected.canRemove)

        settings.codexActiveSource = .managedAccount(id: legacyID)
        settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: snapshot, loadedAt: Date())
        let store = UsageStore(
            fetcher: UsageFetcher(),
            browserDetection: BrowserDetection(cacheTTL: 0),
            settings: settings)
        store.codexAccountSnapshots = [runtimeEnrichedLegacy, siblingProjected].map {
            CodexAccountUsageSnapshot(account: $0, snapshot: nil, error: nil, sourceLabel: "test")
        }
        let controller = StatusItemController(
            store: store,
            settings: settings,
            account: AccountInfo(email: nil, plan: nil),
            updater: DisabledUpdaterController(),
            preferencesSelection: PreferencesSelection(),
            statusBar: .system)
        defer { controller.releaseStatusItemsForTesting() }

        let display = try #require(controller.codexAccountMenuDisplay(for: .codex))

        #expect(legacyProjected.workspaceAccountID == "account-runtime")
        #expect(display.snapshots.map(\.id).sorted() == [legacyProjected.id, siblingProjected.id].sorted())
    }

    @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)
        store._cancelPlanUtilizationHistoryLoadForTesting()
        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()
        }
    }
}