import CodexBarCoreimport Foundationimport Testing private actor AntigravityCredentialUpdateCapture {    private var captured: [AntigravityOAuthCredentials] = []     func append(_ credentials: AntigravityOAuthCredentials) {        self.captured.append(credentials)    }     func values() -> [AntigravityOAuthCredentials] {        self.captured    }} @Suite(.serialized)// swiftlint:disable:next type_body_lengthstruct AntigravityRemoteUsageFetcherTests {    @Test    func `antigravity supports token accounts for quick account switching`() {        let support = TokenAccountSupportCatalog.support(for: .antigravity)         #expect(support?.title == "Google accounts")        #expect(support?.requiresManualCookieSource == false)        #expect(TokenAccountSupportCatalog.envOverride(            for: .antigravity,            token: "serialized-credentials")?[AntigravityOAuthCredentialsStore.environmentCredentialsKey] ==            "serialized-credentials")    }     @Test    func `oauth credentials round trip through token account value`() throws {        let credentials = AntigravityOAuthCredentials(            accessToken: "access-token",            refreshToken: "refresh-token",            expiryDate: Date(timeIntervalSince1970: 1_700_000_000),            idToken: GeminiAPITestHelpers.makeIDToken(email: "user@example.com"),            email: "user@example.com",            projectID: "project-123",            clientID: "client-id",            clientSecret: "client-secret")         let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: credentials)        let decoded = try #require(AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: token))         #expect(decoded == credentials)    }     @Test    func `remote fetch uses selected token account credentials before shared credentials`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "shared-token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "shared@example.com"),            email: "shared@example.com")        let selectedCredentials = AntigravityOAuthCredentials(            accessToken: "selected-token",            refreshToken: nil,            expiryDate: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "selected@example.com"),            email: "selected@example.com",            projectID: nil,            clientID: nil,            clientSecret: nil)        let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: selectedCredentials)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer selected-token")             switch host {            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "currentTier": ["id": "free-tier", "name": "free"],                            "cloudaicompanionProject": "managed-project-123",                        ]))                }                if url.path == "/v1internal:fetchAvailableModels" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: Self.availableModelsResponse())                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let fetcher = AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            environment: [AntigravityOAuthCredentialsStore.environmentCredentialsKey: token],            dataLoader: dataLoader)        let snapshot = try await fetcher.fetch()         #expect(snapshot.accountEmail == "selected@example.com")    }     @Test    func `remote fetch refreshes selected token account without mutating shared credentials`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "shared-token",            refreshToken: "shared-refresh",            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "shared@example.com"),            email: "shared@example.com",            clientID: "shared-client-id",            clientSecret: "shared-client-secret")        let selectedCredentials = AntigravityOAuthCredentials(            accessToken: "selected-old-token",            refreshToken: "selected-refresh",            expiryDate: Date().addingTimeInterval(-3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "selected-old@example.com"),            email: "selected-old@example.com",            projectID: nil,            clientID: "selected-client-id",            clientSecret: "selected-client-secret")        let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: selectedCredentials)        let updateCapture = AntigravityCredentialUpdateCapture()         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }             switch host {            case "oauth2.googleapis.com":                let body = String(data: request.httpBody ?? Data(), encoding: .utf8) ?? ""                #expect(body.contains("client_id=selected-client-id"))                #expect(body.contains("refresh_token=selected-refresh"))                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData([                        "access_token": "selected-new-token",                        "expires_in": 3600,                        "id_token": GeminiAPITestHelpers.makeIDToken(email: "selected-new@example.com"),                    ]))            case "cloudcode-pa.googleapis.com":                #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer selected-new-token")                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "currentTier": ["id": "free-tier", "name": "free"],                            "cloudaicompanionProject": "selected-project-123",                        ]))                }                if url.path == "/v1internal:fetchAvailableModels" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: Self.availableModelsResponse())                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let fetcher = AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            environment: [AntigravityOAuthCredentialsStore.environmentCredentialsKey: token],            dataLoader: dataLoader,            credentialsUpdateHandler: { credentials in                await updateCapture.append(credentials)            })        let snapshot = try await fetcher.fetch()        let shared = try env.readAntigravityCredentials()        let updatedCredentials = await updateCapture.values()         #expect(snapshot.accountEmail == "selected-new@example.com")        #expect(shared["access_token"] as? String == "shared-token")        #expect(shared["email"] as? String == "shared@example.com")        #expect(updatedCredentials.last?.accessToken == "selected-new-token")        #expect(updatedCredentials.last?.projectID == "selected-project-123")    }     @Test    func `remote fetch ignores selected token account project id persistence failure`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        let selectedCredentials = AntigravityOAuthCredentials(            accessToken: "selected-token",            refreshToken: nil,            expiryDate: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "selected@example.com"),            email: "selected@example.com",            projectID: nil,            clientID: nil,            clientSecret: nil)        let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: selectedCredentials)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }             switch host {            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "currentTier": ["id": "free-tier", "name": "free"],                            "cloudaicompanionProject": "selected-project-123",                        ]))                }                if url.path == "/v1internal:fetchAvailableModels" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: Self.availableModelsResponse())                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let fetcher = AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            environment: [AntigravityOAuthCredentialsStore.environmentCredentialsKey: token],            dataLoader: dataLoader,            credentialsUpdateHandler: { _ in                throw CocoaError(.fileWriteUnknown)            })        let snapshot = try await fetcher.fetch()         #expect(snapshot.accountEmail == "selected@example.com")    }     @Test    func `remote fetch rejects invalid selected token account`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "shared-token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "shared@example.com"),            email: "shared@example.com")         let fetcher = AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            environment: [AntigravityOAuthCredentialsStore.environmentCredentialsKey: "not-json"],            dataLoader: GeminiAPITestHelpers.dataLoader { _ in                throw URLError(.badServerResponse)            })         do {            _ = try await fetcher.fetch()            #expect(Bool(false), "Expected selected account decode failure")        } catch let error as AntigravityRemoteFetchError {            guard case let .parseFailed(message) = error else {                #expect(Bool(false), "Unexpected Antigravity error: \(error)")                return            }            #expect(message.contains("selected account"))        } catch {            #expect(Bool(false), "Unexpected error: \(error)")        }    }     @Test    func `remote fetch maps cloud code models into antigravity usage`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "user@company.com", hostedDomain: "company.com"),            email: "user@company.com")         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }             switch host {            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "currentTier": ["id": "standard-tier", "name": "standard"],                            "cloudaicompanionProject": "managed-project-123",                        ]))                }                if url.path == "/v1internal:fetchAvailableModels" {                    let body = try #require(request.httpBody)                    let json = try #require(JSONSerialization.jsonObject(with: body) as? [String: Any])                    #expect(json["project"] as? String == "managed-project-123")                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: Self.availableModelsResponse())                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let fetcher = AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            dataLoader: dataLoader)        let snapshot = try await fetcher.fetch()         #expect(snapshot.accountEmail == "user@company.com")        #expect(snapshot.accountPlan == "Paid")         let usage = try snapshot.toUsageSnapshot()        #expect(usage.primary?.remainingPercent.rounded() == 20)        #expect(usage.secondary?.remainingPercent.rounded() == 50)        #expect(usage.tertiary == nil)    }     @Test    func `remote fetch verifies full model quotas with quota endpoint`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "user@example.com"),            email: "user@example.com")         final class Counter: @unchecked Sendable {            private let lock = NSLock()            private var value = 0             func increment() {                self.lock.lock()                self.value += 1                self.lock.unlock()            }             func get() -> Int {                self.lock.lock()                defer { self.lock.unlock() }                return self.value            }        }         let quotaCalls = Counter()        let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }             switch host {            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistResponse(                            tierId: "standard-tier",                            projectId: "managed-project-123"))                }                if url.path == "/v1internal:fetchAvailableModels" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "models": [                                "claude-sonnet-4": [                                    "displayName": "Claude Sonnet 4",                                    "quotaInfo": ["remainingFraction": 1],                                ],                                "gemini-2.5-pro": [                                    "displayName": "Gemini 2.5 Pro",                                    "quotaInfo": ["remainingFraction": 1],                                ],                                "gemini-2.5-flash": [                                    "displayName": "Gemini 2.5 Flash",                                    "quotaInfo": ["remainingFraction": 1],                                ],                            ],                        ]))                }                if url.path == "/v1internal:retrieveUserQuota" {                    quotaCalls.increment()                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "buckets": [                                [                                    "modelId": "claude-sonnet-4",                                    "resetTime": "2025-01-01T00:00:00Z",                                ],                                [                                    "modelId": "gemini-2.5-pro",                                    "remainingFraction": 0.6,                                    "resetTime": "2025-01-01T00:00:00Z",                                ],                                [                                    "modelId": "gemini-2.5-flash",                                    "remainingFraction": 0.9,                                    "resetTime": "2025-01-01T00:00:00Z",                                ],                            ],                        ]))                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let snapshot = try await AntigravityRemoteUsageFetcher(            timeout: 1,            homeDirectory: env.homeURL.path,            dataLoader: dataLoader)            .fetch()        let usage = try snapshot.toUsageSnapshot()         #expect(quotaCalls.get() == 1)        #expect(usage.primary?.remainingPercent == 60.0)        #expect(usage.secondary?.remainingPercent == 100.0)        #expect(usage.tertiary == nil)    }     @Test    func `remote fetch ignores full model availability when verification has no quota data`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeAntigravityCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: GeminiAPITestHelpers.makeIDToken(email: "user@example.com"),            email: "user@example.com")         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }             switch host {            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistResponse(                            tierId: "standard-tier",                            projectId: "managed-project-123"))                }                if url.path == "/v1internal:fetchAvailableModels" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData([                            "models": [                                "claude-sonnet-4": [                                    "displayName": "Claude Sonnet 4",                                    "quotaInfo": ["remainingFraction": 1],                                ],                                "gemini-2.5-pro": [                                    "displayName": "Gemini 2.5 Pro",                                    "quotaInfo": ["remainingFraction": 1],                                ],                            ],                        ]))                }                if url.path == "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.jsonData(["buckets": []]))                }                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }

Lines 1–500 of 1245. Load more