protocol URLSessionProtocol {
    func data(for request: URLRequest) async throws -> (Data, URLResponse)
    func bytes(for request: URLRequest) async throws -> (URLSession.AsyncBytes, URLResponse)
}

extension URLSession: URLSessionProtocol { }

class MockURLProtocol: URLProtocol {
    typealias RequestHandler: ((URLRequest) throws -> (Data?, URLResponse?, Error?))?
    static var requestHandler: RequestHandler?

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

    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        return request
    }

    override func startLoading() {
        guard let handler = MockURLProtocol.requestHandler else {
            fatalError("Request handler is not set.")
        }

        do {
            let (data, response, error) = try handler(request)
            if let response = response {
                client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
            }
            if let data = data {
                client?.urlProtocol(self, didLoad: data)
            }
            if let error = error {
                client?.urlProtocol(self, didFailWithError: error)
            } else {
                client?.urlProtocolDidFinishLoading(self)
            }
        } catch {
            client?.urlProtocol(self, didFailWithError: error)
        }
    }

    override func stopLoading() {
        // 요청 취소 시 ...
    }
}

final class MockURLSession {
    private let urlSession: URLSession = {
        let configuration = URLSessionConfiguration.ephemeral
        configuration.protocolClasses = [MockURLProtocol.self]
        return URLSession(configuration: configuration)
    }()

    private let isFailRequest: Bool
    private let successMockData: Data

    private let successStatusCode = 200
    private let failureStatusCode = 401

    init(isFailRequest: Bool = false, successMockData: Data) {
        self.isFailRequest = isFailRequest
        self.successMockData = successMockData
    }

    func data(for url: URL) async throws -> (Data, URLResponse) {
        let successResponse = HTTPURLResponse(
            url: url,
            statusCode: successStatusCode,
            httpVersion: "2",
            headerFields: nil
        )
        let failureResponse = HTTPURLResponse(
            url: url,
            statusCode: failureStatusCode,
            httpVersion: "2",
            headerFields: nil
        )

        let response = isFailRequest ? failureResponse : successResponse
        let data = isFailRequest ? nil : successMockData

        MockURLProtocol.requestHandler = { _ in
            return (data, response, nil)
        }

        return try await urlSession.data(for: URLRequest(url: url))
    }
}

let mockData = """
{ "message": "Success" }
""".data(using: .utf8)!

let mockSession = MockURLSession(isFailRequest: false, successMockData: mockData)

let url = URL(string: "https://example.com")!

Task {
    do {
        let (data, response) = try await mockSession.data(for: url)
        if let response = response as? HTTPURLResponse {
            print("Status Code: \(response.statusCode)")
        }
        print("Data: \(String(data: data, encoding: .utf8) ?? "")")
    } catch {
        print("Error: \(error)")
    }
}