import AppKit import CodexBarCore import Foundation import Testing @testable import CodexBar @Suite(.serialized) @MainActor struct StatusMenuCodexSwitcherTests { private func disableMenuCardsForTesting() { StatusItemController.menuCardRenderingEnabled = false StatusItemController.setMenuRefreshEnabledForTesting(false) } private func makeSettings() -> SettingsStore { let suite = "StatusMenuCodexSwitcherTests-\(UUID().uuidString)" let defaults = UserDefaults(suiteName: suite)! defaults.removePersistentDomain(forName: suite) let configStore = testConfigStore(suiteName: suite) return SettingsStore( userDefaults: defaults, configStore: configStore, zaiTokenStore: NoopZaiTokenStore(), syntheticTokenStore: NoopSyntheticTokenStore()) } private func makeStatusBarForTesting() -> NSStatusBar { .system } private func enableOnlyCodex(_ settings: SettingsStore) { let registry = ProviderRegistry.shared for provider in UsageProvider.allCases { guard let metadata = registry.metadata[provider] else { continue } settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex) } } private func makeManagedAccountStoreURL(accounts: [ManagedCodexAccount]) throws -> URL { let storeURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) let store = FileManagedCodexAccountStore(fileURL: storeURL) try store.storeAccounts(ManagedCodexAccountSet( version: FileManagedCodexAccountStore.currentVersion, accounts: accounts)) return storeURL } private func actionLabels(in descriptor: MenuDescriptor) -> [String] { descriptor.sections.flatMap(\.entries).compactMap { entry in guard case let .action(label, _) = entry else { return nil } return label } } private func representedIDs(in menu: NSMenu) -> [String] { menu.items.compactMap { $0.representedObject as? String } } private func snapshot(email: String, percent: Double = 12) -> UsageSnapshot { UsageSnapshot( primary: RateWindow( usedPercent: percent, windowMinutes: 300, resetsAt: Date().addingTimeInterval(300), resetDescription: nil), secondary: RateWindow( usedPercent: percent, windowMinutes: 10080, resetsAt: Date().addingTimeInterval(86400), resetDescription: nil), updatedAt: Date(), identity: ProviderIdentitySnapshot( providerID: .codex, accountEmail: email, accountOrganization: nil, loginMethod: "Plus")) } private func selectCodexVisibleAccountForStatusMenu( id: String, settings: SettingsStore, store: UsageStore) -> Task? { guard settings.selectCodexVisibleAccount(id: id) else { return nil } _ = store.prepareCodexAccountScopedRefreshIfNeeded() return Task { @MainActor in await store.refreshCodexAccountScopedState(allowDisabled: true) } } private func installBlockingCodexProvider(on store: UsageStore, blocker: BlockingStatusMenuCodexFetchStrategy) { let baseSpec = store.providerSpecs[.codex]! store.providerSpecs[.codex] = Self.makeCodexProviderSpec(baseSpec: baseSpec) { try await blocker.awaitResult() } } private static func makeCodexProviderSpec( baseSpec: ProviderSpec, loader: @escaping @Sendable () async throws -> UsageSnapshot) -> ProviderSpec { let baseDescriptor = baseSpec.descriptor let strategy = StatusMenuTestCodexFetchStrategy(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) } @Test func `codex menu shows account switcher and add account action for multiple visible accounts`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let projection = settings.codexVisibleAccountProjection let descriptor = MenuDescriptor.build( provider: .codex, store: store, settings: settings, account: fetcher.loadAccountInfo(), updateReady: false) #expect(projection.visibleAccounts.map(\.email) == ["live@example.com", "managed@example.com"]) #expect(projection.activeVisibleAccountID == "live@example.com") let actionLabels = self.actionLabels(in: descriptor) #expect(actionLabels.contains("Add Account...")) #expect(actionLabels.contains("Switch Account...") == false) } @Test func `codex menu hides account switcher when only one visible account exists`() { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "solo@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) defer { settings._test_liveSystemCodexAccount = nil } let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let descriptor = MenuDescriptor.build( provider: .codex, store: store, settings: settings, account: fetcher.loadAccountInfo(), updateReady: false) #expect(settings.codexVisibleAccountProjection.visibleAccounts.map(\.email) == ["solo@example.com"]) #expect(self.actionLabels(in: descriptor).contains("Add Account...")) } @Test func `codex segmented multi account layout shows account switcher`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false settings.multiAccountMenuLayout = .segmented self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let controller = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu(for: .codex) controller.menuWillOpen(menu) #expect(menu.items.compactMap { $0.view as? CodexAccountSwitcherView }.first != nil) #expect(self.representedIDs(in: menu).filter { $0.hasPrefix("menuCard") } == ["menuCard"]) } @Test func `merged codex menu smart refresh keeps account switcher visible`() throws { self.disableMenuCardsForTesting() StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.setMenuRefreshEnabledForTesting(false) } let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = true settings.selectedMenuProvider = .codex settings.multiAccountMenuLayout = .segmented let registry = ProviderRegistry.shared for provider in UsageProvider.allCases { guard let metadata = registry.metadata[provider] else { continue } settings.setProviderEnabled( provider: provider, metadata: metadata, enabled: provider == .codex || provider == .claude) } let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let controller = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu() controller.menuWillOpen(menu) #expect(menu.items.count(where: { $0.view is CodexAccountSwitcherView }) == 1) controller.menuContentVersion &+= 1 controller.refreshOpenMenusIfNeeded() #expect(menu.items.count(where: { $0.view is CodexAccountSwitcherView }) == 1) settings._test_liveSystemCodexAccount = nil controller.menuContentVersion &+= 1 controller.refreshOpenMenusIfNeeded() #expect(menu.items.count(where: { $0.view is CodexAccountSwitcherView }) == 1) controller.menuDidClose(menu) controller.menuContentVersion &+= 1 controller.menuWillOpen(menu) #expect(menu.items.count(where: { $0.view is CodexAccountSwitcherView }) == 0) } @Test func `codex menu can select preserved switcher row during transient account projection`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .managedAccount(id: managedAccountID) let liveAccount = try #require(settings.codexVisibleAccountProjection.visibleAccounts .first { $0.selectionSource == .liveSystem }) settings._test_liveSystemCodexAccount = nil #expect(settings.codexVisibleAccountProjection.visibleAccounts.map(\.email) == ["managed@example.com"]) settings.selectDisplayedCodexVisibleAccount(liveAccount) #expect(settings.codexActiveSource == .liveSystem) } @Test func `codex stacked multi account layout shows account cards`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false settings.multiAccountMenuLayout = .stacked self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let projection = settings.codexVisibleAccountProjection store.codexAccountSnapshots = projection.visibleAccounts.enumerated().map { index, account in CodexAccountUsageSnapshot( account: account, snapshot: self.snapshot(email: account.email, percent: Double(10 + index)), error: nil, sourceLabel: "test") } let controller = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu(for: .codex) controller.menuWillOpen(menu) #expect(menu.items.compactMap { $0.view as? CodexAccountSwitcherView }.first == nil) #expect(self.representedIDs(in: menu).filter { $0.hasPrefix("menuCard") } == ["menuCard-0", "menuCard-1"]) } @Test func `codex stacked multi account layout shows account cards before per account snapshots load`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false settings.multiAccountMenuLayout = .stacked self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) store._setSnapshotForTesting(self.snapshot(email: "live@example.com"), provider: .codex) let controller = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu(for: .codex) controller.menuWillOpen(menu) #expect(menu.items.compactMap { $0.view as? CodexAccountSwitcherView }.first == nil) #expect(self.representedIDs(in: menu).filter { $0.hasPrefix("menuCard") } == ["menuCard-0", "menuCard-1"]) } @Test func `codex switcher suppresses personal labels while preserving team workspace tooltips`() { let accounts = [ CodexVisibleAccount( id: "live:provider:account-personal", email: "pl.fr@yandex.com", workspaceLabel: "Personal", workspaceAccountID: "account-personal", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed:aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", email: "pl.fr@yandex.com", workspaceLabel: "IDconcepts", workspaceAccountID: "account-team", storedAccountID: UUID(), selectionSource: .managedAccount(id: UUID()), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { _ in }) let titles = view._test_buttonTitles() let toolTips = view._test_buttonToolTips() #expect(titles.count == 2) #expect(titles[0] != titles[1]) #expect(titles.allSatisfy { $0.lowercased().contains("pl.") }) #expect(titles[0].contains("|") == false) #expect(titles[0].lowercased().contains("pers") == false) #expect(titles[1].lowercased().contains("id")) #expect(toolTips == accounts.map(\.menuDisplayName)) #expect(accounts[0].displayName == "pl.fr@yandex.com — Personal") #expect(accounts[0].menuDisplayName == "pl.fr@yandex.com") } @Test func `codex switcher reports fixed menu width for long account labels`() { let accounts = [ CodexVisibleAccount( id: "live:provider:account-personal", email: "managed-account-with-a-very-long-name@example.com", workspaceLabel: nil, workspaceAccountID: "account-managed", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed:aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", email: "steipete-with-a-very-long-label@gmail.com", workspaceLabel: nil, workspaceAccountID: "account-gmail", storedAccountID: UUID(), selectionSource: .managedAccount(id: UUID()), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 310, onSelect: { _ in }) #expect(view.frame.width == 310) #expect(view.intrinsicContentSize.width == 310) #expect(view.fittingSize.width == 310) } @Test func `codex switcher middle truncates long account emails`() { let accounts = [ CodexVisibleAccount( id: "live:provider:account-personal", email: "local-person-with-an-extremely-long-name@example.com", workspaceLabel: nil, workspaceAccountID: "account-managed", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed:aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", email: "second-person-with-an-extremely-long-name@example.com", workspaceLabel: nil, workspaceAccountID: "account-gmail", storedAccountID: UUID(), selectionSource: .managedAccount(id: UUID()), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { _ in }) let titles = view._test_buttonTitles() #expect(titles.count == 2) #expect(titles[0].hasPrefix("local")) #expect(titles[0].contains("…")) #expect(titles[0].hasSuffix(".com")) #expect(titles[1].hasPrefix("second")) #expect(titles[1].contains("…")) #expect(titles[1].hasSuffix(".com")) } @Test func `codex account switcher passes the selected displayed account`() throws { let managedID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let accounts = [ CodexVisibleAccount( id: "live@example.com", email: "live@example.com", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed@example.com", email: "managed@example.com", storedAccountID: managedID, selectionSource: .managedAccount(id: managedID), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] var selectedAccount: CodexVisibleAccount? let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { selectedAccount = $0 }) view._test_selectAccount(id: "managed@example.com") #expect(selectedAccount == accounts[1]) } @Test func `codex menu switcher selection activates the visible managed account`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem #expect(settings.selectCodexVisibleAccount(id: "managed@example.com")) #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) } @Test func `codex menu switcher clears stale account state on the first click`() async throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false settings.costUsageEnabled = false settings.codexCookieSource = .off self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) store._test_codexCreditsLoaderOverride = { CreditsSnapshot(remaining: 0, events: [], updatedAt: Date()) } store._setSnapshotForTesting( UsageSnapshot( primary: RateWindow(usedPercent: 30, windowMinutes: 300, resetsAt: nil, resetDescription: nil), secondary: nil, updatedAt: Date(), identity: ProviderIdentitySnapshot( providerID: .codex, accountEmail: "live@example.com", accountOrganization: nil, loginMethod: "Pro")), provider: .codex) store.lastCodexAccountScopedRefreshGuard = store .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) let blocker = BlockingStatusMenuCodexFetchStrategy() self.installBlockingCodexProvider(on: store, blocker: blocker) let refreshTask = try #require( self.selectCodexVisibleAccountForStatusMenu( id: "managed@example.com", settings: settings, store: store)) await blocker.waitUntilStarted() #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) #expect(store.snapshots[.codex] == nil) await blocker.resume(with: .success( UsageSnapshot( primary: RateWindow(usedPercent: 9, windowMinutes: 300, resetsAt: nil, resetDescription: nil), secondary: nil, updatedAt: Date(), identity: ProviderIdentitySnapshot( providerID: .codex, accountEmail: "managed@example.com", accountOrganization: nil, loginMethod: "Pro")))) for _ in 0..<10 where store.snapshots[.codex]?.accountEmail(for: .codex) != "managed@example.com" { try? await Task.sleep(for: .milliseconds(20)) } await refreshTask.value #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "managed@example.com") } @Test func `codex account state disables add account while managed authentication is in flight`() async throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) defer { settings._test_liveSystemCodexAccount = nil } let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) defer { try? FileManager.default.removeItem(at: root) } let runner = BlockingManagedCodexLoginRunnerForStatusMenuTests() let service = ManagedCodexAccountService( store: InMemoryManagedCodexAccountStoreForStatusMenuTests(), homeFactory: TestManagedCodexHomeFactoryForStatusMenuTests(root: root), loginRunner: runner, identityReader: StubManagedCodexIdentityReaderForStatusMenuTests(email: "managed@example.com"), workspaceResolver: StubManagedCodexWorkspaceResolverForStatusMenuTests()) let coordinator = ManagedCodexAccountCoordinator(service: service) let authTask = Task { try await coordinator.authenticateManagedAccount() } await runner.waitUntilStarted() let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let pane = ProvidersPane( settings: settings, store: store, managedCodexAccountCoordinator: coordinator) let state = try #require(pane._test_codexAccountsSectionState()) #expect(state.canAddAccount == false) #expect(state.isAuthenticatingManagedAccount) #expect(state.addAccountTitle == "Adding Account…") await runner.resume() _ = try await authTask.value } @Test func `codex account state disables add account when managed store is unreadable`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings._test_unreadableManagedCodexAccountStore = true defer { settings._test_liveSystemCodexAccount = nil settings._test_unreadableManagedCodexAccountStore = false } let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) let pane = ProvidersPane(settings: settings, store: store) let state = try #require(pane._test_codexAccountsSectionState()) #expect(state.hasUnreadableManagedAccountStore) #expect(state.canAddAccount == false) } @Test func `codex menu switcher can select managed row when same email rows split by identity`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false self.enableOnlyCodex(settings) let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( UUID().uuidString, isDirectory: true) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-222222222222")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "same@example.com", managedHomePath: managedHome.path, createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) try Self.writeCodexAuthFile( homeURL: managedHome, email: "same@example.com", plan: "pro", accountID: "account-managed") let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) try? FileManager.default.removeItem(at: managedHome) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "SAME@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date(), identity: .emailOnly(normalizedEmail: "same@example.com")) settings.codexActiveSource = .liveSystem let projection = settings.codexVisibleAccountProjection #expect(projection.visibleAccounts.count == 2) let managedVisibleAccount = try #require(projection.visibleAccounts .first { $0.storedAccountID == managedAccountID }) #expect(settings.selectCodexVisibleAccount(id: managedVisibleAccount.id)) #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) } } extension StatusMenuCodexSwitcherTests { @Test func `codex account switcher swallows child button hit testing for first click`() throws { let managedID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let accounts = [ CodexVisibleAccount( id: "live@example.com", email: "live@example.com", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed@example.com", email: "managed@example.com", storedAccountID: managedID, selectionSource: .managedAccount(id: managedID), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { _ in }) #expect(view.acceptsFirstMouse(for: nil) == true) #expect(view._test_hitTestSwallowsChildButton(id: "managed@example.com") == true) #expect(view._test_toolTipAfterHitTest(id: "managed@example.com") == "managed@example.com") } @Test func `codex account switcher routes runtime click path to selected account`() throws { let managedID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let accounts = [ CodexVisibleAccount( id: "live@example.com", email: "live@example.com", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed@example.com", email: "managed@example.com", storedAccountID: managedID, selectionSource: .managedAccount(id: managedID), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), ] var selectedAccount: CodexVisibleAccount? let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { selectedAccount = $0 }) #expect(view._test_simulateRuntimeClick(id: "managed@example.com") == true) #expect(selectedAccount == accounts[1]) } @Test func `codex account switcher runtime click resolves second row buttons`() throws { let managedID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let secondManagedID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-222222222222")) let accounts = [ CodexVisibleAccount( id: "live@example.com", email: "live@example.com", storedAccountID: nil, selectionSource: .liveSystem, isActive: true, isLive: true, canReauthenticate: true, canRemove: false), CodexVisibleAccount( id: "managed@example.com", email: "managed@example.com", storedAccountID: managedID, selectionSource: .managedAccount(id: managedID), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), CodexVisibleAccount( id: "team@example.com", email: "team@example.com", storedAccountID: secondManagedID, selectionSource: .managedAccount(id: secondManagedID), isActive: false, isLive: false, canReauthenticate: true, canRemove: true), CodexVisibleAccount( id: "second-row@example.com", email: "second-row@example.com", storedAccountID: nil, selectionSource: .liveSystem, isActive: false, isLive: true, canReauthenticate: true, canRemove: false), ] var selectedAccount: CodexVisibleAccount? let view = CodexAccountSwitcherView( accounts: accounts, selectedAccountID: accounts.first?.id, width: 220, onSelect: { selectedAccount = $0 }) #expect(view._test_simulateRuntimeClick(id: "second-row@example.com") == true) #expect(selectedAccount == accounts[3]) } } @MainActor extension StatusMenuCodexSwitcherTests { @Test func `codex account switch defers open menu rebuild until after switcher action`() async throws { self.disableMenuCardsForTesting() StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.setMenuRefreshEnabledForTesting(false) } let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = true settings.selectedMenuProvider = .codex settings.multiAccountMenuLayout = .segmented let registry = ProviderRegistry.shared for provider in UsageProvider.allCases { guard let metadata = registry.metadata[provider] else { continue } settings.setProviderEnabled( provider: provider, metadata: metadata, enabled: provider == .codex || provider == .claude) } let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) store._setSnapshotForTesting(self.snapshot(email: "live@example.com", percent: 11), provider: .codex) store.lastCodexAccountScopedRefreshGuard = store.currentCodexAccountScopedRefreshGuard( preferCurrentSnapshot: false) let blocker = BlockingStatusMenuCodexFetchStrategy() self.installBlockingCodexProvider(on: store, blocker: blocker) let controller = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) defer { controller.releaseStatusItemsForTesting() } let menu = controller.makeMenu() controller.menuWillOpen(menu) let switcher = try #require(menu.items.compactMap { $0.view as? CodexAccountSwitcherView }.first) let managedVisibleAccount = try #require(settings.codexVisibleAccountProjection.visibleAccounts .first { $0.storedAccountID == managedAccountID }) var rebuildCount = 0 controller._test_openMenuRebuildObserver = { _ in rebuildCount += 1 } defer { controller._test_openMenuRebuildObserver = nil } switcher._test_selectAccount(id: managedVisibleAccount.id) #expect(rebuildCount == 0) for _ in 0..<20 where rebuildCount == 0 { await Task.yield() } #expect(rebuildCount == 1) await blocker.waitUntilStarted() await blocker.resume(with: .success(self.snapshot(email: "managed@example.com", percent: 17))) } @Test func `codex account scoped refresh does not retain status controller while in flight`() async throws { self.disableMenuCardsForTesting() StatusItemController.setMenuRefreshEnabledForTesting(true) defer { StatusItemController.setMenuRefreshEnabledForTesting(false) } let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = true settings.selectedMenuProvider = .codex settings.multiAccountMenuLayout = .segmented let registry = ProviderRegistry.shared for provider in UsageProvider.allCases { guard let metadata = registry.metadata[provider] else { continue } settings.setProviderEnabled( provider: provider, metadata: metadata, enabled: provider == .codex || provider == .claude) } let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) store._setSnapshotForTesting(self.snapshot(email: "live@example.com", percent: 11), provider: .codex) store.lastCodexAccountScopedRefreshGuard = store.currentCodexAccountScopedRefreshGuard( preferCurrentSnapshot: false) let blocker = BlockingStatusMenuCodexFetchStrategy() self.installBlockingCodexProvider(on: store, blocker: blocker) var controller: StatusItemController? = StatusItemController( store: store, settings: settings, account: fetcher.loadAccountInfo(), updater: DisabledUpdaterController(), preferencesSelection: PreferencesSelection(), statusBar: self.makeStatusBarForTesting()) let releasedController = WeakStatusItemControllerReference(controller) do { let activeController = try #require(controller) let menu = activeController.makeMenu() activeController.menuWillOpen(menu) let switcher = try #require(menu.items.compactMap { $0.view as? CodexAccountSwitcherView }.first) let managedVisibleAccount = try #require(settings.codexVisibleAccountProjection.visibleAccounts .first { $0.storedAccountID == managedAccountID }) switcher._test_selectAccount(id: managedVisibleAccount.id) } await blocker.waitUntilStarted() controller?.releaseStatusItemsForTesting() controller = nil for _ in 0..<20 where releasedController.value != nil { await Task.yield() } #expect(releasedController.value == nil) await blocker.resume(with: .success(self.snapshot(email: "managed@example.com", percent: 17))) } @Test func `codex stacked refresh discards selected outcome when visible selection changes mid flight`() throws { self.disableMenuCardsForTesting() let settings = self.makeSettings() settings.statusChecksEnabled = false settings.refreshFrequency = .manual settings.mergeIcons = false settings.multiAccountMenuLayout = .stacked self.enableOnlyCodex(settings) let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) let managedAccount = ManagedCodexAccount( id: managedAccountID, email: "managed@example.com", managedHomePath: "/tmp/managed-home", createdAt: 1, updatedAt: 2, lastAuthenticatedAt: 2) let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) defer { settings._test_managedCodexAccountStoreURL = nil settings._test_liveSystemCodexAccount = nil try? FileManager.default.removeItem(at: storeURL) } settings._test_managedCodexAccountStoreURL = storeURL settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( email: "live@example.com", codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) store._setSnapshotForTesting(self.snapshot(email: "managed@example.com", percent: 77), provider: .codex) let projection = settings.codexVisibleAccountProjection let originalVisibleAccountID = projection.activeVisibleAccountID let originalSelectionSource = originalVisibleAccountID.flatMap { projection.source(forVisibleAccountID: $0) } #expect(store.codexVisibleSelectionStillMatches( originalVisibleAccountID: originalVisibleAccountID, originalSelectionSource: originalSelectionSource)) #expect(settings.selectCodexVisibleAccount(id: "managed@example.com")) #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) #expect(store.codexVisibleSelectionStillMatches( originalVisibleAccountID: originalVisibleAccountID, originalSelectionSource: originalSelectionSource) == false) #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "managed@example.com") #expect(store.snapshots[.codex]?.primary?.usedPercent == 77) } private static func writeCodexAuthFile( homeURL: URL, email: String, plan: String, accountID: String? = nil) throws { try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) var tokens: [String: Any] = [ "accessToken": "access-token", "refreshToken": "refresh-token", "idToken": Self.fakeJWT(email: email, plan: plan, accountID: accountID), ] if let accountID { tokens["account_id"] = accountID } let auth = ["tokens": tokens] let data = try JSONSerialization.data(withJSONObject: auth) try data.write(to: homeURL.appendingPathComponent("auth.json")) } private static func fakeJWT(email: String, plan: String, accountID: String? = nil) -> String { let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() var payloadObject: [String: Any] = [ "email": email, "chatgpt_plan_type": plan, ] if let accountID { payloadObject["https://api.openai.com/auth"] = [ "chatgpt_account_id": accountID, ] } let payload = (try? JSONSerialization.data(withJSONObject: payloadObject)) ?? Data() func base64URL(_ data: Data) -> String { data.base64EncodedString() .replacingOccurrences(of: "=", with: "") .replacingOccurrences(of: "+", with: "-") .replacingOccurrences(of: "/", with: "_") } return "\(base64URL(header)).\(base64URL(payload))." } } private final class WeakStatusItemControllerReference { weak var value: StatusItemController? init(_ value: StatusItemController?) { self.value = value } } private struct StatusMenuTestCodexFetchStrategy: ProviderFetchStrategy { let loader: @Sendable () async throws -> UsageSnapshot var id: String { "status-menu-test-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: "status-menu-test-codex") } func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { false } } private actor BlockingStatusMenuCodexFetchStrategy { private var waiters: [CheckedContinuation, Never>] = [] private var startedWaiters: [(count: Int, continuation: CheckedContinuation)] = [] private var startCount = 0 func awaitResult() async throws -> UsageSnapshot { let result = await withCheckedContinuation { continuation in self.waiters.append(continuation) self.startCount += 1 self.resumeStartedWaitersIfReady() } return try result.get() } func waitUntilStarted() async { await self.waitForStartCount(1) } func waitForStartCount(_ count: Int) async { if self.startCount >= count { return } await withCheckedContinuation { continuation in self.startedWaiters.append((count, continuation)) } } func resume(with result: Result) { self.waiters.forEach { $0.resume(returning: result) } self.waiters.removeAll() } private func resumeStartedWaitersIfReady() { let readyWaiters = self.startedWaiters.filter { self.startCount >= $0.count } self.startedWaiters.removeAll { self.startCount >= $0.count } readyWaiters.forEach { $0.continuation.resume() } } } private actor BlockingManagedCodexLoginRunnerForStatusMenuTests: ManagedCodexLoginRunning { private var waiters: [CheckedContinuation] = [] private var startedWaiters: [CheckedContinuation] = [] private var didStart = false func run(homePath _: String, timeout _: TimeInterval) async -> CodexLoginRunner.Result { await withCheckedContinuation { continuation in self.waiters.append(continuation) self.didStart = true self.startedWaiters.forEach { $0.resume() } self.startedWaiters.removeAll() } } func waitUntilStarted() async { if self.didStart { return } await withCheckedContinuation { continuation in self.startedWaiters.append(continuation) } } func resume() { let result = CodexLoginRunner.Result(outcome: .success, output: "ok") self.waiters.forEach { $0.resume(returning: result) } self.waiters.removeAll() } } private final class InMemoryManagedCodexAccountStoreForStatusMenuTests: ManagedCodexAccountStoring, @unchecked Sendable { private var snapshot = ManagedCodexAccountSet(version: 1, accounts: []) func loadAccounts() throws -> ManagedCodexAccountSet { self.snapshot } func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { self.snapshot = accounts } func ensureFileExists() throws -> URL { FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) } } private struct StubManagedCodexWorkspaceResolverForStatusMenuTests: ManagedCodexWorkspaceResolving { func resolveWorkspaceIdentity( homePath _: String, providerAccountID _: String) async -> CodexOpenAIWorkspaceIdentity? { nil } } private struct TestManagedCodexHomeFactoryForStatusMenuTests: ManagedCodexHomeProducing { let root: URL func makeHomeURL() -> URL { self.root.appendingPathComponent(UUID().uuidString, isDirectory: true) } func validateManagedHomeForDeletion(_ url: URL) throws { try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) } } private struct StubManagedCodexIdentityReaderForStatusMenuTests: ManagedCodexIdentityReading { let email: String func loadAccountIdentity(homePath _: String) throws -> CodexAuthBackedAccount { CodexAuthBackedAccount( identity: CodexIdentityResolver.resolve(accountId: nil, email: self.email), email: self.email, plan: "Pro") } }