hyungi_document_server/Tests/AITests/MockURLProtocol.swift

import Foundation

/// 테스트용 URLProtocol — URLSession 을 가로채 canned 응답/에러를 돌려주고, 나간 요청을 기록.
/// 라이브 네트워크 0 으로 LocalMLX 의 probe/complete call-shape 를 검증.
final class MockURLProtocol: URLProtocol {
    /// (request) -> (response, body). throw 하면 URLSession 에러 경로.
    nonisolated(unsafe) static var handler: (@Sendable (URLRequest) throws -> (HTTPURLResponse, Data))?
    /// 마지막으로 가로챈 요청(body 포함) 기록.
    nonisolated(unsafe) static var recorder = RequestRecorder()

    static func reset() {
        handler = nil
        recorder = RequestRecorder()
    }

    override class func canInit(with request: URLRequest) -> Bool { true }
    override class func canonicalRequest(for request: URLRequest) -> URLRequest { request }

    override func startLoading() {
        Self.recorder.record(request)
        guard let handler = Self.handler else {
            client?.urlProtocol(self, didFailWithError: URLError(.unsupportedURL))
            return
        }
        do {
            let (response, data) = try handler(request)
            client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
            client?.urlProtocol(self, didLoad: data)
            client?.urlProtocolDidFinishLoading(self)
        } catch {
            client?.urlProtocol(self, didFailWithError: error)
        }
    }

    override func stopLoading() {}

    // MARK: helpers

    static func session() -> URLSession {
        let config = URLSessionConfiguration.ephemeral
        config.protocolClasses = [MockURLProtocol.self]
        return URLSession(configuration: config)
    }

    static func ok(_ url: URL, json: Data) -> (HTTPURLResponse, Data) {
        (HTTPURLResponse(url: url, statusCode: 200, httpVersion: nil, headerFields: nil)!, json)
    }

    static func status(_ url: URL, _ code: Int, body: String = "") -> (HTTPURLResponse, Data) {
        (HTTPURLResponse(url: url, statusCode: code, httpVersion: nil, headerFields: nil)!, Data(body.utf8))
    }
}

/// 나간 요청 기록(body 는 httpBody 또는 httpBodyStream 에서 추출 — URLProtocol 은 보통 stream 으로 전달).
final class RequestRecorder: @unchecked Sendable {
    private(set) var lastURL: URL?
    private(set) var lastMethod: String?
    private(set) var lastBody: Data?
    private(set) var callCount = 0

    func record(_ request: URLRequest) {
        callCount += 1
        lastURL = request.url
        lastMethod = request.httpMethod
        lastBody = request.bodyData
    }
}

extension URLRequest {
    /// URLProtocol 단계에서 body 추출 — httpBody 가 nil 이면 httpBodyStream 에서 읽음.
    var bodyData: Data? {
        if let httpBody { return httpBody }
        guard let stream = httpBodyStream else { return nil }
        stream.open()
        defer { stream.close() }
        var data = Data()
        let bufSize = 8192
        var buffer = [UInt8](repeating: 0, count: bufSize)
        while stream.hasBytesAvailable {
            let read = stream.read(&buffer, maxLength: bufSize)
            if read <= 0 { break }
            data.append(buffer, count: read)
        }
        return data
    }
}