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) }}
import CodexBarCore
import Foundation
import 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)
}
}