405 lines · 14.3 KB · Swift
RawDownload
import Foundationimport Testing@testable import CodexBarCore @Suite(.serialized)struct BedrockUsageStatsTests {    @Test    func `to usage snapshot with budget shows primary window`() {        let snapshot = BedrockUsageSnapshot(            monthlySpend: 50,            monthlyBudget: 200,            inputTokens: 1_500_000,            outputTokens: 500_000,            region: "us-east-1",            updatedAt: Date(timeIntervalSince1970: 1_739_841_600))         let usage = snapshot.toUsageSnapshot()         #expect(usage.primary?.usedPercent == 25)        #expect(usage.primary?.resetDescription == "Monthly budget")        #expect(usage.primary?.resetsAt != nil)        #expect(usage.providerCost?.used == 50)        #expect(usage.providerCost?.limit == 200)        #expect(usage.providerCost?.currencyCode == "USD")        #expect(usage.providerCost?.period == "Monthly")        #expect(usage.identity?.providerID == .bedrock)        #expect(usage.identity?.loginMethod?.contains("Spend: $50.00") == true)    }     @Test    func `to usage snapshot without budget omits primary window`() {        let snapshot = BedrockUsageSnapshot(            monthlySpend: 75.5,            monthlyBudget: nil,            region: "us-west-2",            updatedAt: Date(timeIntervalSince1970: 1_739_841_600))         let usage = snapshot.toUsageSnapshot()         #expect(usage.primary == nil)        #expect(usage.providerCost?.used == 75.5)        #expect(usage.providerCost?.limit == 0)    }     @Test    func `settings reader parses credentials from environment`() {        let env = [            "AWS_ACCESS_KEY_ID": "AKIAIOSFODNN7EXAMPLE",            "AWS_SECRET_ACCESS_KEY": "secret",            "AWS_REGION": "eu-west-1",            "CODEXBAR_BEDROCK_BUDGET": "500",        ]         #expect(BedrockSettingsReader.accessKeyID(environment: env) == "AKIAIOSFODNN7EXAMPLE")        #expect(BedrockSettingsReader.secretAccessKey(environment: env) == "secret")        #expect(BedrockSettingsReader.region(environment: env) == "eu-west-1")        #expect(BedrockSettingsReader.budget(environment: env) == 500)        #expect(BedrockSettingsReader.hasCredentials(environment: env))    }     @Test    func `settings reader requires both credential fields`() {        #expect(!BedrockSettingsReader.hasCredentials(environment: [:]))        #expect(!BedrockSettingsReader.hasCredentials(environment: [            "AWS_ACCESS_KEY_ID": "AKIATEST",        ]))        #expect(!BedrockSettingsReader.hasCredentials(environment: [            "AWS_SECRET_ACCESS_KEY": "secret",        ]))    }     @Test    func `cost explorer response parsing extracts total`() async throws {        let registered = URLProtocol.registerClass(BedrockStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(BedrockStubURLProtocol.self)            }            BedrockStubURLProtocol.handler = nil        }         BedrockStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            let body = """            {                "ResultsByTime": [                    {                        "TimePeriod": {"Start": "2026-04-01", "End": "2026-04-06"},                        "Groups": [                            {                                "Keys": ["Claude Opus (Bedrock Edition)"],                                "Metrics": {"UnblendedCost": {"Amount": "30.00", "Unit": "USD"}}                            },                            {                                "Keys": ["Claude Sonnet (Bedrock Edition)"],                                "Metrics": {"UnblendedCost": {"Amount": "12.50", "Unit": "USD"}}                            },                            {                                "Keys": ["Amazon EC2"],                                "Metrics": {"UnblendedCost": {"Amount": "5.00", "Unit": "USD"}}                            }                        ]                    }                ]            }            """            return Self.makeResponse(url: url, body: body, statusCode: 200)        }         let credentials = BedrockAWSSigner.Credentials(            accessKeyID: "AKIATEST",            secretAccessKey: "testSecret",            sessionToken: nil)         let usage = try await BedrockUsageFetcher.fetchUsage(            credentials: credentials,            region: "us-east-1",            budget: 100,            environment: ["CODEXBAR_BEDROCK_API_URL": "https://bedrock.test"])         #expect(usage.monthlySpend == 42.50)        #expect(usage.monthlyBudget == 100)        #expect(usage.region == "us-east-1")    }     @Test    func `cost explorer data unavailable response returns zero usage`() async throws {        let registered = URLProtocol.registerClass(BedrockStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(BedrockStubURLProtocol.self)            }            BedrockStubURLProtocol.handler = nil        }         BedrockStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            return Self.makeResponse(                url: url,                body: #"{"__type":"com.amazonaws.ce#DataUnavailableException","message":"Data is not ready"}"#,                statusCode: 400)        }         let credentials = BedrockAWSSigner.Credentials(            accessKeyID: "AKIATEST",            secretAccessKey: "testSecret",            sessionToken: nil)         let usage = try await BedrockUsageFetcher.fetchUsage(            credentials: credentials,            region: "us-east-1",            budget: 100,            environment: [BedrockSettingsReader.apiURLKey: "https://bedrock.test"])         #expect(usage.monthlySpend == 0)        #expect(usage.monthlyBudget == 100)    }     @Test    func `cost explorer unrelated bad request remains an API error`() async throws {        let registered = URLProtocol.registerClass(BedrockStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(BedrockStubURLProtocol.self)            }            BedrockStubURLProtocol.handler = nil        }         BedrockStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            return Self.makeResponse(                url: url,                body: #"{"__type":"ValidationException","message":"Invalid request"}"#,                statusCode: 400)        }         let credentials = BedrockAWSSigner.Credentials(            accessKeyID: "AKIATEST",            secretAccessKey: "testSecret",            sessionToken: nil)         await #expect(throws: BedrockUsageError.apiError("HTTP 400")) {            try await BedrockUsageFetcher.fetchUsage(                credentials: credentials,                region: "us-east-1",                budget: nil,                environment: [BedrockSettingsReader.apiURLKey: "https://bedrock.test"])        }    }     @Test    func `cost explorer pagination aggregates monthly total`() async throws {        let registered = URLProtocol.registerClass(BedrockStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(BedrockStubURLProtocol.self)            }            BedrockStubURLProtocol.handler = nil        }         let responses = BedrockStubResponseQueue([            """            {                "NextPageToken": "page-2",                "ResultsByTime": [                    {                        "TimePeriod": {"Start": "2026-04-01", "End": "2026-04-06"},                        "Groups": [                            {                                "Keys": ["Amazon EC2"],                                "Metrics": {"UnblendedCost": {"Amount": "5.00", "Unit": "USD"}}                            }                        ]                    }                ]            }            """,            """            {                "ResultsByTime": [                    {                        "TimePeriod": {"Start": "2026-04-01", "End": "2026-04-06"},                        "Groups": [                            {                                "Keys": ["Amazon Bedrock"],                                "Metrics": {"UnblendedCost": {"Amount": "12.00", "Unit": "USD"}}                            },                            {                                "Keys": ["Claude Sonnet (Bedrock Edition)"],                                "Metrics": {"UnblendedCost": {"Amount": "8.00", "Unit": "USD"}}                            }                        ]                    }                ]            }            """,        ])        BedrockStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            return responses.next(url: url)        }         let credentials = BedrockAWSSigner.Credentials(            accessKeyID: "AKIATEST",            secretAccessKey: "testSecret",            sessionToken: nil)         let usage = try await BedrockUsageFetcher.fetchUsage(            credentials: credentials,            region: "us-east-1",            budget: nil,            environment: [BedrockSettingsReader.apiURLKey: "https://bedrock.test"])         #expect(usage.monthlySpend == 20)        #expect(responses.remainingCount == 0)    }     @Test    func `cost usage fetcher uses provided bedrock environment`() async throws {        let registered = URLProtocol.registerClass(BedrockStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(BedrockStubURLProtocol.self)            }            BedrockStubURLProtocol.handler = nil        }         let responses = BedrockStubResponseQueue([            """            {                "NextPageToken": "daily-page-2",                "ResultsByTime": [                    {                        "TimePeriod": {"Start": "2025-12-10", "End": "2025-12-11"},                        "Groups": [                            {                                "Keys": ["Amazon EC2"],                                "Metrics": {"UnblendedCost": {"Amount": "5.00", "Unit": "USD"}}                            }                        ]                    }                ]            }            """,            """            {                "ResultsByTime": [                    {                        "TimePeriod": {"Start": "2025-12-10", "End": "2025-12-11"},                        "Groups": [                            {                                "Keys": ["Amazon Bedrock"],                                "Metrics": {"UnblendedCost": {"Amount": "7.25", "Unit": "USD"}}                            }                        ]                    }                ]            }            """,        ])        BedrockStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            return responses.next(url: url)        }         let snapshot = try await CostUsageFetcher().loadTokenSnapshot(            provider: .bedrock,            environment: [                BedrockSettingsReader.accessKeyIDKey: "AKIATEST",                BedrockSettingsReader.secretAccessKeyKey: "testSecret",                BedrockSettingsReader.apiURLKey: "https://bedrock.test",            ],            now: Date(timeIntervalSince1970: 1_765_324_800))         #expect(snapshot.last30DaysCostUSD == 7.25)        #expect(snapshot.sessionCostUSD == 7.25)        #expect(snapshot.daily.map(\.date) == ["2025-12-10"])        #expect(responses.remainingCount == 0)    }     @Test    func `current month range uses UTC calendar`() throws {        let originalTimeZone = NSTimeZone.default        NSTimeZone.default = TimeZone(secondsFromGMT: 14 * 60 * 60)!        defer {            NSTimeZone.default = originalTimeZone        }         let now = try #require(ISO8601DateFormatter().date(from: "2026-05-10T12:00:00Z"))        let range = BedrockUsageFetcher.currentMonthRange(now: now)         #expect(range.start == "2026-05-01")        #expect(range.end == "2026-05-11")    }     private final class BedrockStubResponseQueue {        private let lock = NSLock()        private var bodies: [String]         init(_ bodies: [String]) {            self.bodies = bodies        }         var remainingCount: Int {            self.lock.lock()            defer { self.lock.unlock() }            return self.bodies.count        }         func next(url: URL) -> (HTTPURLResponse, Data) {            self.lock.lock()            let body = self.bodies.isEmpty ? #"{"ResultsByTime":[]}"# : self.bodies.removeFirst()            self.lock.unlock()             let response = HTTPURLResponse(                url: url,                statusCode: 200,                httpVersion: "HTTP/1.1",                headerFields: ["Content-Type": "application/json"])!            return (response, Data(body.utf8))        }    }     private static func makeResponse(        url: URL,        body: String,        statusCode: Int = 200) -> (HTTPURLResponse, Data)    {        let response = HTTPURLResponse(            url: url,            statusCode: statusCode,            httpVersion: "HTTP/1.1",            headerFields: ["Content-Type": "application/json"])!        return (response, Data(body.utf8))    }} final class BedrockStubURLProtocol: URLProtocol {    nonisolated(unsafe) static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))?     override static func canInit(with request: URLRequest) -> Bool {        request.url?.host == "bedrock.test"    }     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() {}}