874 lines
32 KiB
Swift
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() {}
|
|
}
|