spb/zyquo-router Public MIT

1//2//  GatewayBehaviorTests.swift3//  Zyquo Router4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//8//  Phase 7.4 — gateway behavior against a mock OpenAI-compatible upstream9//  (built from the same HTTPServer): client-disconnect cancels the upstream10//  stream, fallback chains report the actually-used model, transient 429s11//  retry, and graceful shutdown ends an active stream.12//1314import XCTest15@testable import ZyquoRouter1617final class GatewayBehaviorTests: XCTestCase {18    private let upstreamPort = 1890119    private let routerPort = 189022021    /// Signals observed inside the mock upstream.22    private actor UpstreamState {23        var attempts = 024        var streamCancelled = false2526        func recordAttempt() -> Int {27            attempts += 128            return attempts29        }3031        func markCancelled() {32            streamCancelled = true33        }34    }3536    private func model(_ id: String, provider: ProviderID) -> AIModel {37        AIModel(38            id: id, provider: provider, displayName: id,39            contextWindow: 128_000, maxOutputTokens: 512,40            capabilities: ModelCapabilities(tools: true),41            pricing: nil,42            parameterSupport: ParameterSupport(),43            customBaseURL: URL(string: "http://127.0.0.1:\(upstreamPort)/v1")44        )45    }4647    /// Mock upstream: model "slow" streams forever (marks cancellation when48    /// the client goes away), "fail-500" always 500s, "flaky-429" 429s once49    /// then succeeds, "good" answers immediately.50    private func startMockUpstream(state: UpstreamState) async -> Task<Void, Error> {51        let handler: @Sendable (RouteRequest) async -> RouteResult = { request in52            let body = (try? JSONSerialization.jsonObject(with: request.body)) as? [String: Any] ?? [:]53            let model = body["model"] as? String ?? ""54            let stream = body["stream"] as? Bool ?? false5556            func completion(_ text: String) -> Data {57                ChunkEmitter.serialize([58                    "id": "mock-1", "object": "chat.completion", "created": 1, "model": model,59                    "choices": [["index": 0, "message": ["role": "assistant", "content": text],60                                 "finish_reason": "stop"] as [String: Any]],61                    "usage": ["prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2],62                ])63            }6465            switch model {66            case "fail-500":67                return .complete(status: .internalServerError, headers: [], body: Data("{\"error\":{\"message\":\"boom\"}}".utf8))68            case "flaky-429":69                let attempt = await state.recordAttempt()70                if attempt == 1 {71                    return .complete(status: .tooManyRequests, headers: [("Retry-After", "0")], body: Data("{}".utf8))72                }73                return .complete(status: .ok, headers: [("Content-Type", "application/json")], body: completion("second try"))74            case "slow" where stream:75                return .stream(status: .ok, headers: []) { writer in76                    do {77                        for index in 0..<600 {78                            try await writer.send(raw: ChunkEmitter.serialize([79                                "id": "mock-s", "object": "chat.completion.chunk", "created": 1, "model": model,80                                "choices": [["index": 0, "delta": ["content": "tick\(index) "],81                                             "finish_reason": NSNull()] as [String: Any]],82                            ]))83                            try await Task.sleep(nanoseconds: 50_000_000)84                        }85                    } catch {86                        await state.markCancelled()87                        throw error88                    }89                }90            default:91                return .complete(status: .ok, headers: [("Content-Type", "application/json")], body: completion("hello from good"))92            }93        }9495        let server = HTTPServer(host: "127.0.0.1", port: upstreamPort, handler: handler)96        let started = expectation(description: "mock upstream up")97        let task = Task { try await server.run { started.fulfill() } }98        await fulfillment(of: [started], timeout: 5)99        return task100    }101102    private func startRouter(catalog: [AIModel], chains: [String: [String]] = [:]) async -> Task<Void, Error> {103        let routes = Routes(104            router: RequestRouter(catalog: catalog, fallbackChains: chains),105            providerKey: { _ in "mock-key" }106        )107        let server = HTTPServer(host: "127.0.0.1", port: routerPort) { request in108            await routes.handle(request)109        }110        let started = expectation(description: "router up")111        let task = Task { try await server.run { started.fulfill() } }112        await fulfillment(of: [started], timeout: 5)113        return task114    }115116    private func post(_ body: [String: Any]) async throws -> (Int, [String: Any]) {117        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)118        request.httpMethod = "POST"119        request.setValue("application/json", forHTTPHeaderField: "Content-Type")120        request.httpBody = try JSONSerialization.data(withJSONObject: body)121        let (data, response) = try await URLSession.shared.data(for: request)122        let json = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] ?? [:]123        return ((response as! HTTPURLResponse).statusCode, json)124    }125126    func testClientDisconnectCancelsUpstream() async throws {127        let state = UpstreamState()128        let upstream = await startMockUpstream(state: state)129        let router = await startRouter(catalog: [model("slow", provider: .together)])130        defer { upstream.cancel(); router.cancel() }131132        // Open a streaming request, read a couple of chunks, then drop it.133        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)134        request.httpMethod = "POST"135        request.setValue("application/json", forHTTPHeaderField: "Content-Type")136        request.httpBody = try JSONSerialization.data(withJSONObject: [137            "model": "together/slow",138            "messages": [["role": "user", "content": "go"]],139            "stream": true,140        ])141        let client = Task {142            let (bytes, _) = try await URLSession.shared.bytes(for: request)143            var seen = 0144            for try await _ in bytes.lines {145                seen += 1146                if seen >= 4 { break } // abandon mid-stream147            }148        }149        _ = try? await client.value150151        // The router must cancel its upstream call promptly after the client152        // vanishes (write failure propagates → upstream stream torn down).153        let deadline = Date().addingTimeInterval(8)154        while Date() < deadline {155            if await state.streamCancelled { break }156            try await Task.sleep(nanoseconds: 100_000_000)157        }158        let cancelled = await state.streamCancelled159        XCTAssertTrue(cancelled, "upstream stream was not cancelled after client disconnect")160    }161162    func testFallbackChainReportsActuallyUsedModel() async throws {163        let state = UpstreamState()164        let upstream = await startMockUpstream(state: state)165        let router = await startRouter(166            catalog: [model("fail-500", provider: .together), model("good", provider: .deepinfra)],167            chains: ["together/fail-500": ["deepinfra/good"]]168        )169        defer { upstream.cancel(); router.cancel() }170171        let (status, json) = try await post([172            "model": "together/fail-500",173            "messages": [["role": "user", "content": "go"]],174        ])175        XCTAssertEqual(status, 200)176        XCTAssertEqual(json["model"] as? String, "deepinfra/good", "must report the model that answered")177    }178179    func testRetryOnTransient429() async throws {180        let state = UpstreamState()181        let upstream = await startMockUpstream(state: state)182        let router = await startRouter(catalog: [model("flaky-429", provider: .together)])183        defer { upstream.cancel(); router.cancel() }184185        let (status, json) = try await post([186            "model": "together/flaky-429",187            "messages": [["role": "user", "content": "go"]],188        ])189        XCTAssertEqual(status, 200, "transient 429 must be retried: \(json)")190        let attempts = await state.attempts191        XCTAssertEqual(attempts, 2)192        let content = (((json["choices"] as? [[String: Any]])?.first?["message"] as? [String: Any])?["content"] as? String)193        XCTAssertEqual(content, "second try")194    }195196    func testGracefulShutdownWithActiveStream() async throws {197        let state = UpstreamState()198        let upstream = await startMockUpstream(state: state)199        let router = await startRouter(catalog: [model("slow", provider: .together)])200        defer { upstream.cancel() }201202        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)203        request.httpMethod = "POST"204        request.setValue("application/json", forHTTPHeaderField: "Content-Type")205        request.httpBody = try JSONSerialization.data(withJSONObject: [206            "model": "together/slow",207            "messages": [["role": "user", "content": "go"]],208            "stream": true,209        ])210        let client = Task { () -> Int in211            let (bytes, _) = try await URLSession.shared.bytes(for: request)212            var seen = 0213            for try await _ in bytes.lines { seen += 1 }214            return seen215        }216217        // Give the stream time to start flowing, then stop the router.218        try await Task.sleep(nanoseconds: 700_000_000)219        router.cancel()220        _ = try? await router.value221222        // The client's connection must terminate (not hang) once the server223        // shuts down, and the port must be immediately rebindable.224        let seen = (try? await client.value) ?? -1225        XCTAssertNotEqual(seen, -1, "client saw some chunks then a clean termination")226227        let reborn = await startRouter(catalog: [model("good", provider: .together)])228        reborn.cancel()229        _ = try? await reborn.value230    }231}232