Files
steipete--codexbar/Tests/CodexBarTests/OpenCodeGoUsageFetcherErrorTests.swift
2026-07-13 12:22:33 +08:00

874 lines
32 KiB
Swift

import Foundation
import Testing
@testable import CodexBarCore
private final class OpenCodeGoRequestRecorder<Value: Sendable>: @unchecked Sendable {
private let lock = NSLock()
private var storage: [Value] = []
func append(_ value: Value) {
self.lock.lock()
defer { self.lock.unlock() }
self.storage.append(value)
}
var values: [Value] {
self.lock.lock()
defer { self.lock.unlock() }
return self.storage
}
}
private final class OpenCodeGoContinuationBox<Value: Sendable>: @unchecked Sendable {
private let lock = NSLock()
private var continuation: CheckedContinuation<Value, Never>?
func wait(onReady: @Sendable () -> Void) async -> Value {
await withCheckedContinuation { continuation in
self.lock.lock()
self.continuation = continuation
self.lock.unlock()
onReady()
}
}
func resume(returning value: Value) {
self.lock.lock()
let continuation = self.continuation
self.continuation = nil
self.lock.unlock()
continuation?.resume(returning: value)
}
}
@Suite(.serialized)
struct OpenCodeGoUsageFetcherErrorTests {
@Test
func `dashboard URL uses normalized workspace ID`() {
#expect(
OpenCodeGoUsageFetcher.dashboardURL(workspaceID: "https://opencode.ai/workspace/wrk_abc123/go")
.absoluteString == "https://opencode.ai/workspace/wrk_abc123/go")
#expect(
OpenCodeGoUsageFetcher.dashboardURL(workspaceID: "workspace=wrk_def456")
.absoluteString == "https://opencode.ai/workspace/wrk_def456/go")
#expect(
OpenCodeGoUsageFetcher.dashboardURL(workspaceID: nil)
.absoluteString == "https://opencode.ai")
}
private struct UsageWindow {
let percent: Double
let resetInSec: Int
}
private func makeSession() -> URLSession {
let config = URLSessionConfiguration.ephemeral
config.protocolClasses = [OpenCodeGoStubURLProtocol.self]
return URLSession(configuration: config)
}
@Test
func `redirect guard allows only same-host https redirects`() {
#expect(OpenCodeGoUsageFetcher.allowsRedirect(
from: URL(string: "https://opencode.ai/_server"),
to: URL(string: "https://opencode.ai/workspace/wrk_TEST123/go")))
#expect(!OpenCodeGoUsageFetcher.allowsRedirect(
from: URL(string: "https://opencode.ai/_server"),
to: URL(string: "https://evil.example/steal")))
#expect(!OpenCodeGoUsageFetcher.allowsRedirect(
from: URL(string: "https://opencode.ai/_server"),
to: URL(string: "http://opencode.ai/insecure")))
}
@Test
func `extracts api error from detail field`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
let body = #"{"detail":"Workspace missing"}"#
return Self.makeResponse(url: url, body: body, statusCode: 500, contentType: "application/json")
}
do {
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
Issue.record("Expected OpenCodeGoUsageError.apiError")
} catch let error as OpenCodeGoUsageError {
switch error {
case let .apiError(message):
#expect(message.contains("HTTP 500"))
#expect(message.contains("Workspace missing"))
default:
Issue.record("Expected apiError, got: \(error)")
}
}
}
@Test
func `workspace get missing ids falls back to post before loading go page`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let requests = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
requests.append("\(request.httpMethod ?? "GET") \(url.path)")
let workspaceServerID = "def39973159c7f0483d8793a822b8dbb10d067e12c65455fcb4608459ba0234f"
if url.query?.contains(workspaceServerID) == true,
request.httpMethod?.uppercased() == "GET"
{
return Self.makeResponse(
url: url,
body: #"{"ok":true}"#,
statusCode: 200,
contentType: "application/json")
}
if url.path == "/_server",
request.httpMethod?.uppercased() == "POST",
request.value(forHTTPHeaderField: "X-Server-Id") == workspaceServerID
{
return Self.makeResponse(
url: url,
body: #"{"data":[{"id":"wrk_TEST123"}]}"#,
statusCode: 200,
contentType: "application/json")
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 22, resetInSec: 300),
weekly: UsageWindow(percent: 44, resetInSec: 3600),
monthly: UsageWindow(percent: 55, resetInSec: 7200)),
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
includeZenBalance: false,
session: self.makeSession())
#expect(snapshot.rollingUsagePercent == 22)
#expect(snapshot.weeklyUsagePercent == 44)
#expect(snapshot.monthlyUsagePercent == 55)
#expect(requests.values == [
"GET /_server",
"POST /_server",
"GET /workspace/wrk_TEST123/go",
])
}
@Test
func `workspace get public actor error is treated as invalid credentials without post retry`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let methods = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
methods.append(request.httpMethod ?? "GET")
let body = [
#";0x00000263;((self.$R=self.$R||{})["server-fn:test"]=[],"#,
#"($R=>$R[0]=Object.assign(new Error("actor of type \"public\" is not associated with an account"),"#,
#"{stack:"Error: actor of type \"public\" is not associated with an account"}))"#,
#"($R["server-fn:test"]))"#,
].joined()
return Self.makeResponse(
url: url,
body: body,
statusCode: 200,
contentType: "text/javascript")
}
do {
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
session: self.makeSession())
Issue.record("Expected OpenCodeGoUsageError.invalidCredentials")
} catch let error as OpenCodeGoUsageError {
switch error {
case .invalidCredentials:
break
default:
Issue.record("Expected invalidCredentials, got: \(error)")
}
}
#expect(methods.values == ["GET"])
}
@Test
func `go page missing usage fields returns parse failed without post retry`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let methods = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
methods.append(request.httpMethod ?? "GET")
return Self.makeResponse(
url: url,
body: "<html><title>opencode</title><body>No usage yet</body></html>",
statusCode: 200,
contentType: "text/html")
}
do {
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
Issue.record("Expected OpenCodeGoUsageError.parseFailed")
} catch let error as OpenCodeGoUsageError {
switch error {
case let .parseFailed(message):
#expect(message.contains("Missing usage fields"))
default:
Issue.record("Expected parseFailed, got: \(error)")
}
}
#expect(methods.values == ["GET", "GET", "GET"])
}
@Test
func `zen only account waits for balance beyond optional grace`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let observedPaths = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedPaths.append(url.path)
if url.path == "/workspace/wrk_TEST123" {
Thread.sleep(forTimeInterval: 0.4)
return Self.makeResponse(
url: url,
body: #"<html><body><h2>現在の残高 $42.50</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: "<html><title>opencode</title><body>No Go subscription usage</body></html>",
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
let usage = snapshot.toUsageSnapshot()
#expect(snapshot.isBalanceOnly)
#expect(snapshot.zenBalanceUSD == 42.5)
#expect(usage.primary == nil)
#expect(usage.secondary == nil)
#expect(usage.providerCost?.used == 42.5)
#expect(usage.providerCost?.period == "Zen balance")
#expect(observedPaths.values.count == 2)
#expect(Set(observedPaths.values) == ["/workspace/wrk_TEST123/go", "/workspace/wrk_TEST123"])
}
@Test
func `zen only account fetches required balance when optional usage is disabled`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let observedPaths = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedPaths.append(url.path)
if url.path == "/workspace/wrk_TEST123" {
return Self.makeResponse(
url: url,
body: #"<html><body><h2>Current balance $23.75</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: "<html><title>opencode</title><body>No Go subscription usage</body></html>",
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
includeZenBalance: false,
session: self.makeSession())
#expect(snapshot.isBalanceOnly)
#expect(snapshot.zenBalanceUSD == 23.75)
#expect(observedPaths.values == ["/workspace/wrk_TEST123/go", "/workspace/wrk_TEST123"])
}
@Test
func `zen only account propagates invalid credentials from required balance fetch`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
return Self.makeResponse(
url: url,
body: "Unauthorized",
statusCode: 401,
contentType: "text/plain")
}
return Self.makeResponse(
url: url,
body: "<html><title>opencode</title><body>No Go subscription usage</body></html>",
statusCode: 200,
contentType: "text/html")
}
do {
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=stale",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
Issue.record("Expected invalid credentials to propagate.")
} catch OpenCodeGoUsageError.invalidCredentials {
// Expected.
} catch {
Issue.record("Expected invalidCredentials, got: \(error)")
}
}
@Test
func `zen only account falls back after final subscription parse failure`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
var rootTimeout: TimeInterval?
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
rootTimeout = request.timeoutInterval
return Self.makeResponse(
url: url,
body: #"<html><body><h2>Current balance $17.25</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: #"<script>rollingUsage:{usagePercent:12}</script>"#,
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 12,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
let usage = snapshot.toUsageSnapshot()
#expect(snapshot.isBalanceOnly)
#expect(snapshot.zenBalanceUSD == 17.25)
#expect(usage.primary == nil)
#expect(usage.secondary == nil)
#expect(usage.providerCost?.used == 17.25)
#expect(rootTimeout == 12)
}
@Test
func `zen only fallback promptly cancels when balance task ignores cancellation`() async throws {
let balanceStarted = AsyncStream<Void>.makeStream(of: Void.self)
let balanceContinuation = OpenCodeGoContinuationBox<Double?>()
let balanceTask = Task<Double?, Error> {
await balanceContinuation.wait {
balanceStarted.continuation.yield(())
}
}
let fallbackTask = Task {
try await OpenCodeGoUsageFetcher.requiredZenBalanceFallback(
from: balanceTask,
for: .parseFailed("Missing usage fields."),
request: OpenCodeGoUsageFetcher.ZenBalanceRequest(
workspaceID: "wrk_TEST123",
cookieHeader: "auth=test",
timeout: 2,
session: self.makeSession()),
now: Date())
}
var iterator = balanceStarted.stream.makeAsyncIterator()
_ = await iterator.next()
let start = ContinuousClock.now
fallbackTask.cancel()
do {
_ = try await fallbackTask.value
Issue.record("Expected cancellation to propagate.")
} catch is CancellationError {
// Expected.
} catch {
Issue.record("Expected CancellationError, got: \(error)")
}
#expect(start.duration(to: .now) < .milliseconds(500))
#expect(balanceTask.isCancelled)
balanceContinuation.resume(returning: 42.5)
#expect(try await balanceTask.value == 42.5)
}
@Test
func `normalizes workspace override from URL into go page path`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let observedPaths = OpenCodeGoRequestRecorder<String>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedPaths.append(url.path)
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_URL123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "https://opencode.ai/workspace/wrk_URL123/billing",
session: self.makeSession())
#expect(observedPaths.values.count == 3)
#expect(Set(observedPaths.values) == [
"/workspace/wrk_URL123/go",
"/workspace/wrk_URL123",
"/_server",
])
}
@Test
func `fetcher attaches optional zen balance from workspace root`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
return Self.makeResponse(
url: url,
body: #"<html><body><h2>現在の残高 $98.76</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
#expect(snapshot.zenBalanceUSD == 98.76)
#expect(snapshot.toUsageSnapshot().providerCost?.period == "Zen balance")
}
@Test
func `fetcher falls back to billing server when workspace page omits balance`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let observedRequests = OpenCodeGoRequestRecorder<URLRequest>()
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedRequests.append(request)
if url.path == "/workspace/wrk_TEST123" {
return Self.makeResponse(
url: url,
body: "<html><body>Workspace dashboard without hydrated billing data</body></html>",
statusCode: 200,
contentType: "text/html")
}
if url.path == "/_server" {
return Self.makeResponse(
url: url,
body: #"$R[0]={customerID:"cus_test",balance:$R[1]=9876000000,reload:!1}"#,
statusCode: 200,
contentType: "text/javascript")
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
#expect(snapshot.zenBalanceUSD == 98.76)
let billingRequest = try #require(observedRequests.values.first { $0.url?.path == "/_server" })
let billingURL = try #require(billingRequest.url)
let components = try #require(URLComponents(url: billingURL, resolvingAgainstBaseURL: false))
let query = Dictionary(uniqueKeysWithValues: (components.queryItems ?? []).map { ($0.name, $0.value ?? "") })
#expect(query["id"] == "c83b78a614689c38ebee981f9b39a8b377716db85c1fd7dbab604adc02d3313d")
#expect(query["args"] == #"["wrk_TEST123"]"#)
#expect(billingRequest.value(forHTTPHeaderField: "Cookie") == "auth=test")
#expect(billingRequest.value(forHTTPHeaderField: "Referer") == "https://opencode.ai/workspace/wrk_TEST123")
}
@Test
func `optional zen balance helper uses normalized cookie and workspace override`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
var observedCookie: String?
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedCookie = request.value(forHTTPHeaderField: "Cookie")
#expect(url.path == "/workspace/wrk_TEST123")
return Self.makeResponse(
url: url,
body: #"<html><body><h2>現在の残高 $98.76</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
let balance = try await OpenCodeGoUsageFetcher.fetchOptionalZenBalance(
cookieHeader: "provider=google; auth=test",
timeout: 2,
workspaceIDOverride: "https://opencode.ai/workspace/wrk_TEST123/go",
session: self.makeSession())
#expect(balance == 98.76)
#expect(observedCookie == "auth=test")
}
@Test
func `optional zen balance failure does not fail subscription usage`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
var rootTimeout: TimeInterval?
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
rootTimeout = request.timeoutInterval
throw URLError(.timedOut)
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 60,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
#expect(snapshot.rollingUsagePercent == 17)
#expect(snapshot.zenBalanceUSD == nil)
#expect(rootTimeout == 60)
}
@Test
func `optional zen balance does not stall subscription usage`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
Thread.sleep(forTimeInterval: 1)
return Self.makeResponse(
url: url,
body: #"<html><body><h2>現在の残高 $98.76</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let start = ContinuousClock.now
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 60,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
let elapsed = start.duration(to: ContinuousClock.now)
#expect(snapshot.rollingUsagePercent == 17)
#expect(snapshot.zenBalanceUSD == nil)
#expect(elapsed < .milliseconds(700))
}
@Test
func `optional zen balance can be skipped by settings`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
var observedPaths: [String] = []
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedPaths.append(url.path)
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 60,
workspaceIDOverride: "wrk_TEST123",
includeZenBalance: false,
session: self.makeSession())
#expect(snapshot.rollingUsagePercent == 17)
#expect(snapshot.zenBalanceUSD == nil)
#expect(observedPaths == ["/workspace/wrk_TEST123/go"])
}
@Test
func `optional zen balance cancellation propagates`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
let rootStarted = AsyncStream<Void>.makeStream(of: Void.self)
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
if url.path == "/workspace/wrk_TEST123" {
rootStarted.continuation.yield(())
Thread.sleep(forTimeInterval: 0.2)
return Self.makeResponse(
url: url,
body: #"<html><body><h2>現在の残高 $98.76</h2></body></html>"#,
statusCode: 200,
contentType: "text/html")
}
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
let task = Task {
try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "auth=test",
timeout: 60,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
}
let started = await withTaskGroup(of: Bool.self) { group in
group.addTask {
var iterator = rootStarted.stream.makeAsyncIterator()
return await iterator.next() != nil
}
group.addTask {
try? await Task.sleep(for: .seconds(2))
return false
}
let result = await group.next() ?? false
group.cancelAll()
return result
}
#expect(started)
task.cancel()
do {
_ = try await task.value
Issue.record("Expected cancellation to propagate.")
} catch is CancellationError {
// Expected.
} catch {
Issue.record("Expected CancellationError, got: \(error)")
}
}
@Test
func `fetcher sends only auth cookie to opencode host`() async throws {
defer {
OpenCodeGoStubURLProtocol.handler = nil
}
var observedCookie: String?
OpenCodeGoStubURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
observedCookie = request.value(forHTTPHeaderField: "Cookie")
return Self.makeResponse(
url: url,
body: Self.goUsagePageHTML(
workspaceID: "wrk_TEST123",
rolling: UsageWindow(percent: 17, resetInSec: 600),
weekly: UsageWindow(percent: 75, resetInSec: 7200),
monthly: nil),
statusCode: 200,
contentType: "text/html")
}
_ = try await OpenCodeGoUsageFetcher.fetchUsage(
cookieHeader: "provider=google; auth=test",
timeout: 2,
workspaceIDOverride: "wrk_TEST123",
session: self.makeSession())
#expect(observedCookie == "auth=test")
}
private static func goUsagePageHTML(
workspaceID: String,
rolling: UsageWindow,
weekly: UsageWindow,
monthly: UsageWindow?) -> String
{
let monthlyField: String? = if let monthly {
#"monthlyUsage:{status:"ok",resetInSec:\#(monthly.resetInSec),usagePercent:\#(monthly.percent)}"#
} else {
nil
}
let usageFields = [
#"rollingUsage:{status:"ok",resetInSec:\#(rolling.resetInSec),usagePercent:\#(rolling.percent)}"#,
#"weeklyUsage:{status:"ok",resetInSec:\#(weekly.resetInSec),usagePercent:\#(weekly.percent)}"#,
monthlyField,
]
.compactMap(\.self)
.joined(separator: ",")
return """
<!DOCTYPE html>
<html>
<body>
<script>
_$HY.r["lite.subscription.get[\\"\(workspaceID)\\"]"]=$R[17]=$R[2]($R[18]={p:0,s:0,f:0});
$R[24]($R[18],$R[27]={mine:!0,useBalance:!1,\(usageFields)});
</script>
</body>
</html>
"""
}
private static func makeResponse(
url: URL,
body: String,
statusCode: Int,
contentType: String) -> (HTTPURLResponse, Data)
{
let response = HTTPURLResponse(
url: url,
statusCode: statusCode,
httpVersion: "HTTP/1.1",
headerFields: ["Content-Type": contentType])!
return (response, Data(body.utf8))
}
}
final class OpenCodeGoStubURLProtocol: URLProtocol {
nonisolated(unsafe) static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))?
override static func canInit(with request: URLRequest) -> Bool {
request.url?.host == "opencode.ai"
}
override static func canonicalRequest(for request: URLRequest) -> URLRequest {
request
}
override func startLoading() {
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() {}
}