import Foundation import Testing @testable import CodexBarCore @Suite(.serialized) struct CodebuffUsageFetcherTests { @Test func `usage URL composes the correct endpoint`() throws { let base = try #require(URL(string: "https://www.codebuff.com")) let url = CodebuffUsageFetcher.usageURL(baseURL: base) #expect(url.absoluteString == "https://www.codebuff.com/api/v1/usage") } @Test func `subscription URL composes the correct endpoint`() throws { let base = try #require(URL(string: "https://www.codebuff.com")) let url = CodebuffUsageFetcher.subscriptionURL(baseURL: base) #expect(url.absoluteString == "https://www.codebuff.com/api/user/subscription") } @Test func `usage request sends required fingerprint id`() async throws { defer { CodebuffStubURLProtocol.handler = nil CodebuffStubURLProtocol.requests = [] CodebuffStubURLProtocol.requestBodies = [] } CodebuffStubURLProtocol.requests = [] CodebuffStubURLProtocol.requestBodies = [] CodebuffStubURLProtocol.handler = { request in guard let url = request.url else { throw URLError(.badURL) } switch url.path { case "/api/v1/usage": return try Self.makeResponse(url: url, body: #"{"usage":25,"quota":100,"remainingBalance":75}"#) case "/api/user/subscription": return try Self.makeResponse(url: url, body: "{}") default: return try Self.makeResponse(url: url, body: "{}", statusCode: 404) } } let snapshot = try await CodebuffUsageFetcher.fetchUsage(apiKey: "cb-test", session: Self.makeSession()) let usageIndex = try #require(CodebuffStubURLProtocol.requests.firstIndex { $0.url?.path == "/api/v1/usage" }) let body = try #require(CodebuffStubURLProtocol.requestBodies[usageIndex]) let payload = try #require(JSONSerialization.jsonObject(with: body) as? [String: String]) #expect(payload["fingerprintId"] == "codexbar-usage") #expect(snapshot.creditsUsed == 25) } @Test func `usage fetch can skip subscription endpoint for API key tokens`() async throws { defer { CodebuffStubURLProtocol.handler = nil CodebuffStubURLProtocol.requests = [] CodebuffStubURLProtocol.requestBodies = [] } CodebuffStubURLProtocol.requests = [] CodebuffStubURLProtocol.requestBodies = [] CodebuffStubURLProtocol.handler = { request in guard let url = request.url else { throw URLError(.badURL) } switch url.path { case "/api/v1/usage": return try Self.makeResponse(url: url, body: #"{"usage":25,"quota":100,"remainingBalance":75}"#) case "/api/user/subscription": Issue.record("Subscription endpoint should not be called for API key tokens") return try Self.makeResponse(url: url, body: "{}", statusCode: 500) default: return try Self.makeResponse(url: url, body: "{}", statusCode: 404) } } let snapshot = try await CodebuffUsageFetcher.fetchUsage( apiKey: "cb-test", includeSubscription: false, session: Self.makeSession()) #expect(snapshot.creditsUsed == 25) #expect(CodebuffStubURLProtocol.requests.map(\.url?.path) == ["/api/v1/usage"]) } @Test func `subscription grace does not wait for transport that ignores cancellation`() async throws { let transport = ProviderHTTPTransportStub { request in let url = try #require(request.url) if url.path == "/api/v1/usage" { let response = try Self.makeResponse( url: url, body: #"{"usage":25,"quota":100,"remainingBalance":75}"#) return (response.1, response.0) } let response = try Self.makeResponse( url: url, body: #"{"subscription":{"displayName":"Pro","status":"active"}}"#) return await withCheckedContinuation { continuation in DispatchQueue.global().asyncAfter(deadline: .now() + 0.5) { continuation.resume(returning: (response.1, response.0)) } } } let startedAt = ContinuousClock.now let snapshot = try await CodebuffUsageFetcher._fetchUsageForTesting( apiKey: "cb-test", transport: transport, subscriptionGrace: .milliseconds(20)) let elapsed = startedAt.duration(to: .now) #expect(snapshot.creditsUsed == 25) #expect(snapshot.tier == nil) #expect(elapsed < .milliseconds(300), "Subscription enrichment delayed usage: \(elapsed)") // Let the deliberately cancellation-ignoring test task drain before the test exits. try await Task.sleep(for: .milliseconds(550)) } @Test func `cancellation stops subscription while usage transport ignores cancellation`() async throws { let usageStarted = CodebuffRequestGate() let subscriptionStarted = CodebuffRequestGate() let subscriptionCancelled = CodebuffRequestGate() let transport = ProviderHTTPTransportStub { request in let url = try #require(request.url) if url.path == "/api/v1/usage" { await usageStarted.open() let response = try Self.makeResponse( url: url, body: #"{"usage":25,"quota":100,"remainingBalance":75}"#) return await withCheckedContinuation { continuation in DispatchQueue.global().asyncAfter(deadline: .now() + 0.5) { continuation.resume(returning: (response.1, response.0)) } } } await subscriptionStarted.open() do { try await Task.sleep(for: .seconds(10)) } catch { await subscriptionCancelled.open() throw error } let response = try Self.makeResponse( url: url, body: #"{"subscription":{"displayName":"Pro","status":"active"}}"#) return (response.1, response.0) } let task = Task { try await CodebuffUsageFetcher.fetchUsage( apiKey: "cb-test", session: transport) } await usageStarted.wait() await subscriptionStarted.wait() let cancellationStartedAt = ContinuousClock.now task.cancel() await subscriptionCancelled.wait() #expect(cancellationStartedAt.duration(to: .now) < .milliseconds(300)) await #expect(throws: CancellationError.self) { try await task.value } } @Test func `api strategy only fetches subscription for credentials file tokens`() { let envResolution = ProviderTokenResolution(token: "env-token", source: .environment) let fileResolution = ProviderTokenResolution(token: "file-token", source: .authFile) #expect(CodebuffAPIFetchStrategy.shouldFetchSubscription(for: envResolution) == false) #expect(CodebuffAPIFetchStrategy.shouldFetchSubscription(for: fileResolution) == true) } @Test func `status 401 maps to unauthorized`() { #expect(CodebuffUsageFetcher._statusErrorForTesting(401) == .unauthorized) #expect(CodebuffUsageFetcher._statusErrorForTesting(403) == .unauthorized) } @Test func `status 404 maps to endpoint not found`() { #expect(CodebuffUsageFetcher._statusErrorForTesting(404) == .endpointNotFound) } @Test func `status 500 maps to service unavailable`() { guard case .serviceUnavailable(503) = CodebuffUsageFetcher._statusErrorForTesting(503) else { Issue.record("Expected .serviceUnavailable(503)") return } } @Test func `status 200 returns nil`() { #expect(CodebuffUsageFetcher._statusErrorForTesting(200) == nil) } @Test func `usage payload parses numeric credit fields`() throws { let json = """ { "usage": 1250, "quota": 5000, "remainingBalance": 3750, "autoTopupEnabled": true, "next_quota_reset": "2026-05-01T00:00:00Z" } """ let payload = try CodebuffUsageFetcher._parseUsagePayloadForTesting(Data(json.utf8)) #expect(payload.used == 1250) #expect(payload.total == 5000) #expect(payload.remaining == 3750) #expect(payload.autoTopupEnabled == true) #expect(payload.nextQuotaReset != nil) } @Test func `usage payload accepts string-encoded numbers`() throws { let json = """ { "usage": "12", "quota": "100", "remainingBalance": "88" } """ let payload = try CodebuffUsageFetcher._parseUsagePayloadForTesting(Data(json.utf8)) #expect(payload.used == 12) #expect(payload.total == 100) #expect(payload.remaining == 88) } @Test func `usage payload returns nil fields when absent`() throws { let payload = try CodebuffUsageFetcher._parseUsagePayloadForTesting(Data("{}".utf8)) #expect(payload.used == nil) #expect(payload.total == nil) #expect(payload.remaining == nil) #expect(payload.autoTopupEnabled == nil) } @Test func `usage payload throws on malformed JSON`() { #expect { _ = try CodebuffUsageFetcher._parseUsagePayloadForTesting(Data("not-json".utf8)) } throws: { error in guard case CodebuffUsageError.parseFailed = error else { return false } return true } } @Test func `subscription payload parses tier and weekly window`() throws { let json = """ { "hasSubscription": true, "subscription": { "status": "active", "tier": "pro", "billingPeriodEnd": "2026-05-15T00:00:00Z" }, "rateLimit": { "weeklyUsed": 2100, "weeklyLimit": 7000, "weeklyResetsAt": "2026-05-08T00:00:00Z" }, "email": "user@example.com" } """ let payload = try CodebuffUsageFetcher._parseSubscriptionPayloadForTesting(Data(json.utf8)) #expect(payload.tier == "pro") #expect(payload.status == "active") #expect(payload.weeklyUsed == 2100) #expect(payload.weeklyLimit == 7000) #expect(payload.weeklyResetsAt != nil) #expect(payload.email == "user@example.com") #expect(payload.billingPeriodEnd != nil) } @Test func `subscription payload prefers display name over numeric tier`() throws { let json = """ { "subscription": { "tier": 2, "displayName": "Pro" } } """ let payload = try CodebuffUsageFetcher._parseSubscriptionPayloadForTesting(Data(json.utf8)) #expect(payload.tier == "Pro") } @Test func `subscription payload falls back to numeric scheduled tier`() throws { let json = """ { "subscription": { "scheduledTier": 3 } } """ let payload = try CodebuffUsageFetcher._parseSubscriptionPayloadForTesting(Data(json.utf8)) #expect(payload.tier == "3") } @Test func `subscription payload formats oversized numeric tier without trapping`() throws { let json = """ { "subscription": { "scheduledTier": 9223372036854775808 } } """ let payload = try CodebuffUsageFetcher._parseSubscriptionPayloadForTesting(Data(json.utf8)) #expect(payload.tier == "9223372036854775808") } @Test func `subscription payload tolerates missing rate limit`() throws { let json = """ { "subscription": { "status": "trialing", "tier": "free" } } """ let payload = try CodebuffUsageFetcher._parseSubscriptionPayloadForTesting(Data(json.utf8)) #expect(payload.weeklyUsed == nil) #expect(payload.weeklyLimit == nil) #expect(payload.status == "trialing") } @Test func `snapshot maps to rate window with credits window`() { let snapshot = CodebuffUsageSnapshot( creditsUsed: 250, creditsTotal: 1000, creditsRemaining: 750, weeklyUsed: 100, weeklyLimit: 500, weeklyResetsAt: Date(timeIntervalSince1970: 1_777_680_000), tier: "pro", autoTopUpEnabled: true, updatedAt: Date()) let unified = snapshot.toUsageSnapshot() #expect(unified.primary?.usedPercent == 25) // The credit balance is intentionally NOT stored in `resetDescription` — // generic renderers prepend "Resets " when `resetsAt` is absent, which would // surface misleading text like "Resets 250/1,000 credits". #expect(unified.primary?.resetDescription == nil) #expect(unified.secondary?.usedPercent == 20) #expect(unified.secondary?.windowMinutes == 7 * 24 * 60) #expect(unified.secondary?.resetsAt == Date(timeIntervalSince1970: 1_777_680_000)) #expect(unified.secondary?.resetDescription == nil) #expect(unified.identity?.providerID == .codebuff) #expect(unified.identity?.loginMethod?.contains("Pro") == true) #expect(unified.identity?.loginMethod?.contains("auto top-up") == true) } @Test func `snapshot infers total from used plus remaining`() { let snapshot = CodebuffUsageSnapshot( creditsUsed: 40, creditsTotal: nil, creditsRemaining: 60) let unified = snapshot.toUsageSnapshot() #expect(unified.primary?.usedPercent == 40) } @Test func `snapshot surfaces exhausted state when quota is missing from payload`() { // Only `creditsUsed` is populated (no total, no remaining) — the API response is // degenerate but we still want the row to be visible so the user notices the // missing configuration instead of seeing an empty/healthy-looking bar. let usedOnly = CodebuffUsageSnapshot( creditsUsed: 42, creditsTotal: nil, creditsRemaining: nil) #expect(usedOnly.toUsageSnapshot().primary?.usedPercent == 100) // Only `creditsRemaining` is populated — same fallback should apply. let remainingOnly = CodebuffUsageSnapshot( creditsUsed: nil, creditsTotal: nil, creditsRemaining: 17) #expect(remainingOnly.toUsageSnapshot().primary?.usedPercent == 100) } @Test func `snapshot hides credit window when no credit fields are present`() { let empty = CodebuffUsageSnapshot() #expect(empty.toUsageSnapshot().primary == nil) } @Test func `missing credentials fetch call throws missing credentials`() async { do { _ = try await CodebuffUsageFetcher.fetchUsage(apiKey: " ") Issue.record("Expected missingCredentials error") } catch let error as CodebuffUsageError { #expect(error == .missingCredentials) } catch { Issue.record("Unexpected error: \(error)") } } private static func makeSession() -> URLSession { let config = URLSessionConfiguration.ephemeral config.protocolClasses = [CodebuffStubURLProtocol.self] return URLSession(configuration: config) } private static func makeResponse( url: URL, body: String, statusCode: Int = 200) throws -> (HTTPURLResponse, Data) { guard let response = HTTPURLResponse( url: url, statusCode: statusCode, httpVersion: nil, headerFields: ["Content-Type": "application/json"]) else { throw URLError(.badServerResponse) } return (response, Data(body.utf8)) } } private actor CodebuffRequestGate { private var isOpen = false private var continuations: [CheckedContinuation] = [] func wait() async { guard !self.isOpen else { return } await withCheckedContinuation { continuation in self.continuations.append(continuation) } } func open() { self.isOpen = true let continuations = self.continuations self.continuations.removeAll() for continuation in continuations { continuation.resume() } } } final class CodebuffStubURLProtocol: URLProtocol { nonisolated(unsafe) static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))? nonisolated(unsafe) static var requests: [URLRequest] = [] nonisolated(unsafe) static var requestBodies: [Data?] = [] override static func canInit(with request: URLRequest) -> Bool { request.url?.host == "www.codebuff.com" } override static func canonicalRequest(for request: URLRequest) -> URLRequest { request } override func startLoading() { Self.requests.append(self.request) Self.requestBodies.append(Self.bodyData(from: self.request)) guard let handler = Self.handler else { self.client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse)) return } do { let (response, data) = try handler(self.request) self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) self.client?.urlProtocol(self, didLoad: data) self.client?.urlProtocolDidFinishLoading(self) } catch { self.client?.urlProtocol(self, didFailWithError: error) } } override func stopLoading() {} private static func bodyData(from request: URLRequest) -> Data? { if let httpBody = request.httpBody { return httpBody } guard let stream = request.httpBodyStream else { return nil } stream.open() defer { stream.close() } var data = Data() var buffer = [UInt8](repeating: 0, count: 1024) while stream.hasBytesAvailable { let count = stream.read(&buffer, maxLength: buffer.count) if count > 0 { data.append(buffer, count: count) } else { break } } return data.isEmpty ? nil : data } }