387 lines · 16.9 KB · Swift
RawDownload
import CodexBarCoreimport Foundationimport Testing @Suite(.serialized)struct GeminiStatusProbePlanTests {    @Test    func `selects project id for quota requests`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let expectedProject = "gen-lang-client-123"        let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                let json = GeminiAPITestHelpers.jsonData([                    "projects": [                        ["projectId": expectedProject],                    ],                ])                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 200, body: json)            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistFreeTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                let bodyText = request.httpBody.flatMap { String(data: $0, encoding: .utf8) } ?? ""                if !bodyText.contains(expectedProject) {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 400, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleFlashQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.modelQuotas.contains { $0.percentLeft == 40 })        let usage = snapshot.toUsageSnapshot()        #expect(usage.secondary?.remainingPercent == 40)        #expect(usage.tertiary == nil)    }     @Test    func `prefers load code assist project for quota requests`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let loadCodeAssistProject = "cloudaicompanion-123"        let fallbackProject = "gen-lang-client-should-not-use"        let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                let json = GeminiAPITestHelpers.jsonData([                    "projects": [                        ["projectId": fallbackProject],                    ],                ])                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 200, body: json)            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistResponse(                            tierId: "free-tier",                            projectId: loadCodeAssistProject))                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                let bodyText = request.httpBody.flatMap { String(data: $0, encoding: .utf8) } ?? ""                if !bodyText.contains(loadCodeAssistProject) {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 400, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleFlashQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.modelQuotas.contains { $0.percentLeft == 40 })        let usage = snapshot.toUsageSnapshot()        #expect(usage.secondary?.remainingPercent == 40)        #expect(usage.tertiary == nil)    }     @Test    func `separates flash and flash lite quota buckets from api response`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistStandardTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        let usage = snapshot.toUsageSnapshot()         #expect(usage.primary?.remainingPercent == 60.0)        #expect(usage.secondary?.remainingPercent == 90.0)        #expect(usage.tertiary?.remainingPercent == 80.0)    }     @Test    func `detects paid from standard tier`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistStandardTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.accountPlan == "Paid")    }     @Test    func `detects workspace from free tier with hosted domain`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        let idToken = GeminiAPITestHelpers.makeIDToken(email: "user@company.com", hostedDomain: "company.com")        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: idToken)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistFreeTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.accountPlan == "Workspace")    }     @Test    func `detects free from free tier without hosted domain`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        let idToken = GeminiAPITestHelpers.makeIDToken(email: "user@gmail.com")        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: idToken)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistFreeTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleFlashQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.accountPlan == "Free")    }     @Test    func `detects legacy from legacy tier`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 200,                        body: GeminiAPITestHelpers.loadCodeAssistLegacyTierResponse())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.accountPlan == "Legacy")    }     @Test    func `leaves blank when load code assist fails`() async throws {        let env = try GeminiTestEnvironment()        defer { env.cleanup() }        try env.writeCredentials(            accessToken: "token",            refreshToken: nil,            expiry: Date().addingTimeInterval(3600),            idToken: nil)         let dataLoader = GeminiAPITestHelpers.dataLoader { request in            guard let url = request.url, let host = url.host else {                throw URLError(.badURL)            }            switch host {            case "cloudresourcemanager.googleapis.com":                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.jsonData(["projects": []]))            case "cloudcode-pa.googleapis.com":                if url.path == "/v1internal:loadCodeAssist" {                    return GeminiAPITestHelpers.response(                        url: url.absoluteString,                        status: 500,                        body: Data())                }                if url.path != "/v1internal:retrieveUserQuota" {                    return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())                }                return GeminiAPITestHelpers.response(                    url: url.absoluteString,                    status: 200,                    body: GeminiAPITestHelpers.sampleQuotaResponse())            default:                return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())            }        }         let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)        let snapshot = try await probe.fetch()        #expect(snapshot.accountPlan == nil)    }}