181 lines · 7.6 KB · Swift
RawDownload
import CodexBarCoreimport Foundationimport Testing@testable import CodexBar @MainActorstruct MenuDescriptorCodexManagedFallbackTests {    @Test    func `codex account section prefers managed fallback over ambient account`() throws {        let suite = "MenuDescriptorCodexManagedFallbackTests"        let defaults = try #require(UserDefaults(suiteName: suite))        defaults.removePersistentDomain(forName: suite)         let settings = SettingsStore(            userDefaults: defaults,            configStore: testConfigStore(suiteName: suite),            zaiTokenStore: NoopZaiTokenStore(),            syntheticTokenStore: NoopSyntheticTokenStore(),            codexCookieStore: InMemoryCookieHeaderStore(),            claudeCookieStore: InMemoryCookieHeaderStore(),            cursorCookieStore: InMemoryCookieHeaderStore(),            opencodeCookieStore: InMemoryCookieHeaderStore(),            factoryCookieStore: InMemoryCookieHeaderStore(),            minimaxCookieStore: InMemoryMiniMaxCookieStore(),            minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(),            kimiTokenStore: InMemoryKimiTokenStore(),            kimiK2TokenStore: InMemoryKimiK2TokenStore(),            augmentCookieStore: InMemoryCookieHeaderStore(),            ampCookieStore: InMemoryCookieHeaderStore(),            copilotTokenStore: InMemoryCopilotTokenStore(),            tokenAccountStore: InMemoryTokenAccountStore())        settings.statusChecksEnabled = false         let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent(            UUID().uuidString,            isDirectory: true)        let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent(            UUID().uuidString,            isDirectory: true)        defer {            try? FileManager.default.removeItem(at: ambientHome)            try? FileManager.default.removeItem(at: managedHome)        }         try Self.writeCodexAuthFile(homeURL: ambientHome, email: "ambient@example.com", plan: "plus")        try Self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "enterprise")        settings.codexActiveSource = .managedAccount(id: UUID())        settings._test_activeManagedCodexRemoteHomePath = managedHome.path         let fetcher = UsageFetcher(environment: ["CODEX_HOME": ambientHome.path])        let store = UsageStore(            fetcher: fetcher,            browserDetection: BrowserDetection(cacheTTL: 0),            settings: settings,            startupBehavior: .testing)        store._setSnapshotForTesting(            UsageSnapshot(                primary: RateWindow(usedPercent: 10, windowMinutes: 300, resetsAt: nil, resetDescription: nil),                secondary: RateWindow(usedPercent: 20, windowMinutes: 10080, resetsAt: nil, resetDescription: nil),                updatedAt: Date(),                identity: nil),            provider: .codex)         let descriptor = MenuDescriptor.build(            provider: .codex,            store: store,            settings: settings,            account: fetcher.loadAccountInfo(),            updateReady: false,            includeContextualActions: false)         let lines = descriptor.sections            .flatMap(\.entries)            .compactMap { entry -> String? in                guard case let .text(text, _) = entry else { return nil }                return text            }         #expect(lines.contains("Account: managed@example.com"))        #expect(lines.contains("Plan: Enterprise"))        #expect(!lines.contains("Account: ambient@example.com"))        #expect(!lines.contains("Plan: Plus"))    }     @Test    func `codex weekly only window renders without session row`() throws {        let suite = "MenuDescriptorCodexManagedFallbackTests-weekly-only"        let defaults = try #require(UserDefaults(suiteName: suite))        defaults.removePersistentDomain(forName: suite)         let settings = SettingsStore(            userDefaults: defaults,            configStore: testConfigStore(suiteName: suite),            zaiTokenStore: NoopZaiTokenStore(),            syntheticTokenStore: NoopSyntheticTokenStore(),            codexCookieStore: InMemoryCookieHeaderStore(),            claudeCookieStore: InMemoryCookieHeaderStore(),            cursorCookieStore: InMemoryCookieHeaderStore(),            opencodeCookieStore: InMemoryCookieHeaderStore(),            factoryCookieStore: InMemoryCookieHeaderStore(),            minimaxCookieStore: InMemoryMiniMaxCookieStore(),            minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(),            kimiTokenStore: InMemoryKimiTokenStore(),            kimiK2TokenStore: InMemoryKimiK2TokenStore(),            augmentCookieStore: InMemoryCookieHeaderStore(),            ampCookieStore: InMemoryCookieHeaderStore(),            copilotTokenStore: InMemoryCopilotTokenStore(),            tokenAccountStore: InMemoryTokenAccountStore())        settings.statusChecksEnabled = false         let fetcher = UsageFetcher()        let store = UsageStore(            fetcher: fetcher,            browserDetection: BrowserDetection(cacheTTL: 0),            settings: settings,            startupBehavior: .testing)        store._setSnapshotForTesting(            UsageSnapshot(                primary: nil,                secondary: RateWindow(                    usedPercent: 20,                    windowMinutes: 10080,                    resetsAt: nil,                    resetDescription: "Apr 6, 2026"),                updatedAt: Date(),                identity: ProviderIdentitySnapshot(                    providerID: .codex,                    accountEmail: "ratulanimation@gmail.com",                    accountOrganization: nil,                    loginMethod: "free")),            provider: .codex)         let descriptor = MenuDescriptor.build(            provider: .codex,            store: store,            settings: settings,            account: fetcher.loadAccountInfo(),            updateReady: false,            includeContextualActions: false)         let lines = descriptor.sections            .flatMap(\.entries)            .compactMap { entry -> String? in                guard case let .text(text, _) = entry else { return nil }                return text            }         #expect(!lines.contains(where: { $0.hasPrefix("Session:") }))        #expect(lines.contains(where: { $0.hasPrefix("Weekly:") }))    }     private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws {        try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true)        let auth = [            "tokens": [                "accessToken": "access-token",                "refreshToken": "refresh-token",                "idToken": Self.fakeJWT(email: email, plan: plan),            ],        ]        let data = try JSONSerialization.data(withJSONObject: auth)        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,        ])) ?? Data()         func base64URL(_ data: Data) -> String {            data.base64EncodedString()                .replacingOccurrences(of: "=", with: "")                .replacingOccurrences(of: "+", with: "-")                .replacingOccurrences(of: "/", with: "_")        }         return "\(base64URL(header)).\(base64URL(payload))."    }}