Files
2026-07-13 12:22:33 +08:00

271 lines
11 KiB
Swift

import CodexBarCore
import Foundation
import Testing
struct GeminiStatusProbeTests {
/// Sample /stats output from Gemini CLI (actual format with box-drawing chars)
static let sampleStatsOutput = """
│ Model Usage Reqs Usage left │
│ ─────────────────────────────────────────────────────────────────────── │
│ gemini-2.5-flash - 99.8% (Resets in 20h 37m) │
│ gemini-2.5-flash-lite - 99.8% (Resets in 20h 37m) │
│ gemini-2.5-pro - 100.0% (Resets in 24h) │
│ gemini-3-pro-preview - 100.0% (Resets in 24h) │
"""
// MARK: - Legacy CLI parsing tests (kept for fallback support)
@Test
func `parses minimum percent from multiple models`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
#expect(snap.dailyPercentLeft == 99.8)
#expect(snap.resetDescription == "Resets in 20h 37m")
}
@Test
func `parses lower percent correctly`() throws {
let output = """
│ Model Usage Reqs Usage left │
│ gemini-2.5-flash 10 85.5% (Resets in 12h) │
│ gemini-2.5-pro 5 92.0% (Resets in 12h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
#expect(snap.dailyPercentLeft == 85.5)
#expect(snap.resetDescription == "Resets in 12h")
}
@Test
func `handles zero percent usage`() throws {
let output = """
│ gemini-2.5-flash 50 0.0% (Resets in 6h) │
│ gemini-2.5-pro 20 15.0% (Resets in 6h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
#expect(snap.dailyPercentLeft == 0.0)
#expect(snap.resetDescription == "Resets in 6h")
}
@Test
func `handles100 percent remaining`() throws {
let output = """
│ gemini-2.5-flash - 100.0% (Resets in 24h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
#expect(snap.dailyPercentLeft == 100.0)
}
@Test
func `throws on empty output`() {
#expect(throws: GeminiStatusProbeError.self) {
try GeminiStatusProbe.parse(text: "")
}
}
@Test
func `throws on no usage data`() {
let output = """
Welcome to Gemini CLI!
Type /help for available commands.
"""
#expect(throws: GeminiStatusProbeError.self) {
try GeminiStatusProbe.parse(text: output)
}
}
@Test
func `strips ANSI codes before parsing`() throws {
let output =
"\u{1B}[32m│\u{1B}[0m gemini-2.5-flash - 75.5% " +
"(Resets in 18h) │"
let snap = try GeminiStatusProbe.parse(text: output)
#expect(snap.dailyPercentLeft == 75.5)
}
@Test
func `preserves raw text`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
#expect(snap.rawText == Self.sampleStatsOutput)
#expect(snap.accountEmail == nil) // Legacy parse doesn't extract email
#expect(snap.accountPlan == nil)
}
@Test
func `parses various reset descriptions`() throws {
let cases: [(String, String)] = [
("Resets in 24h", "Resets in 24h"),
("Resets in 1h 30m", "Resets in 1h 30m"),
("Resets tomorrow", "Resets tomorrow"),
]
for (resetStr, expected) in cases {
let output = """
│ gemini-2.5-flash - 50.0% (\(resetStr)) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
#expect(snap.resetDescription == expected)
}
}
@Test
func `throws not logged in on auth prompt`() {
let authOutputs = [
"Waiting for auth... (Press ESC or CTRL+C to cancel)",
"Login with Google\nUse Gemini API key",
"Some preamble\nWaiting for auth\nMore text",
]
for output in authOutputs {
#expect(throws: GeminiStatusProbeError.notLoggedIn) {
try GeminiStatusProbe.parse(text: output)
}
}
}
// MARK: - Model quota grouping tests
@Test
func `parses models into quota array`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
// Should parse multiple models (exact count may change as Google adds/removes models)
#expect(snap.modelQuotas.count >= 2)
// All model IDs should start with "gemini"
for quota in snap.modelQuotas {
#expect(quota.modelId.hasPrefix("gemini"))
}
}
@Test
func `lowest percent left returns minimum`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
// Flash models are 99.8%, Pro models are 100%, so min should be 99.8
#expect(snap.lowestPercentLeft == 99.8)
}
@Test
func `tier grouping by keyword`() throws {
// Test that flash/pro keyword filtering works (model names may change)
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
let flashQuotas = snap.modelQuotas.filter { $0.modelId.contains("flash") && !$0.modelId.contains("flash-lite") }
let flashLiteQuotas = snap.modelQuotas.filter { $0.modelId.contains("flash-lite") }
let proQuotas = snap.modelQuotas.filter { $0.modelId.contains("pro") }
// Should have at least one of each tier in sample
#expect(!flashQuotas.isEmpty)
#expect(!flashLiteQuotas.isEmpty)
#expect(!proQuotas.isEmpty)
}
@Test
func `tier minimum calculation`() throws {
// Use controlled test data to verify min-per-tier logic
// Model names must start with "gemini-" to match the parser regex
let output = """
│ gemini-a-flash 10 85.0% (Resets in 24h) │
│ gemini-b-flash 5 92.0% (Resets in 24h) │
│ gemini-c-pro 2 95.0% (Resets in 24h) │
│ gemini-d-pro 1 99.0% (Resets in 24h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
let flashQuotas = snap.modelQuotas.filter { $0.modelId.contains("flash") }
let proQuotas = snap.modelQuotas.filter { $0.modelId.contains("pro") }
let flashMin = flashQuotas.min(by: { $0.percentLeft < $1.percentLeft })
let proMin = proQuotas.min(by: { $0.percentLeft < $1.percentLeft })
#expect(flashMin?.percentLeft == 85.0)
#expect(proMin?.percentLeft == 95.0)
}
@Test
func `quotas have reset descriptions`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
// At least some quotas should have reset descriptions
let hasResets = snap.modelQuotas.contains { $0.resetDescription != nil }
#expect(hasResets)
}
@Test
func `to usage snapshot creates separate flash and flash lite meters`() throws {
let snap = try GeminiStatusProbe.parse(text: Self.sampleStatsOutput)
let usage = snap.toUsageSnapshot()
#expect(usage.primary?.remainingPercent == 100.0)
#expect(usage.secondary?.remainingPercent == 99.8)
#expect(usage.tertiary?.remainingPercent == 99.8)
#expect(usage.primary?.windowMinutes == 1440)
#expect(usage.secondary?.windowMinutes == 1440)
#expect(usage.tertiary?.windowMinutes == 1440)
#expect(usage.secondary?.resetDescription == "Resets in 20h 37m")
#expect(usage.tertiary?.resetDescription == "Resets in 20h 37m")
}
@Test
func `to usage snapshot does not let flash lite contaminate flash bucket`() throws {
let output = """
│ gemini-2.5-flash 10 91.0% (Resets in 12h) │
│ gemini-2.5-flash-lite 5 33.0% (Resets in 6h) │
│ gemini-2.5-pro 2 80.0% (Resets in 24h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
let usage = snap.toUsageSnapshot()
#expect(usage.secondary?.remainingPercent == 91.0)
#expect(usage.tertiary?.remainingPercent == 33.0)
#expect(usage.secondary?.resetDescription == "Resets in 12h")
#expect(usage.tertiary?.resetDescription == "Resets in 6h")
}
@Test
func `to usage snapshot omits tertiary when no flash lite exists`() throws {
let output = """
│ gemini-2.5-flash 10 85.0% (Resets in 12h) │
│ gemini-2.5-pro 2 95.0% (Resets in 24h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
let usage = snap.toUsageSnapshot()
#expect(usage.secondary?.remainingPercent == 85.0)
#expect(usage.tertiary == nil)
}
@Test
func `to usage snapshot uses lowest remaining quota per tier`() throws {
let output = """
│ gemini-a-flash 10 91.0% (Resets in 12h) │
│ gemini-b-flash 5 74.0% (Resets in 10h) │
│ gemini-c-flash-lite 8 66.0% (Resets in 9h) │
│ gemini-d-flash-lite 1 42.0% (Resets in 7h) │
│ gemini-e-pro 2 88.0% (Resets in 24h) │
│ gemini-f-pro 1 97.0% (Resets in 25h) │
"""
let snap = try GeminiStatusProbe.parse(text: output)
let usage = snap.toUsageSnapshot()
#expect(usage.primary?.remainingPercent == 88.0)
#expect(usage.secondary?.remainingPercent == 74.0)
#expect(usage.tertiary?.remainingPercent == 42.0)
#expect(usage.secondary?.resetDescription == "Resets in 10h")
#expect(usage.tertiary?.resetDescription == "Resets in 7h")
}
// MARK: - Live API test
@Test
func `live gemini fetch`() async throws {
guard ProcessInfo.processInfo.environment["LIVE_GEMINI_FETCH"] == "1" else {
return
}
let probe = GeminiStatusProbe()
let snap = try await probe.fetch()
print(
"""
Live Gemini usage (via API):
models: \(snap.modelQuotas.map { "\($0.modelId): \($0.percentLeft)%" }.joined(separator: ", "))
lowest: \(snap.lowestPercentLeft ?? -1)% left
""")
#expect(!snap.modelQuotas.isEmpty)
}
}