diff --git a/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse.swift b/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse.swift index ee7f11592..24b194b10 100644 --- a/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse.swift +++ b/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse.swift @@ -138,6 +138,17 @@ extension HTTPClientResponse { } } +@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) +extension HTTPClientResponse { + /// Response body as `ByteBuffer`. + /// - Parameter maxBytes: The maximum number of bytes this method is allowed to accumulate. + /// - Returns: Bytes collected over time + public func bytes(upTo maxBytes: Int) async throws -> ByteBuffer { + let expectedBytes = self.headers.first(name: "content-length").flatMap(Int.init) ?? maxBytes + return try await self.body.collect(upTo: expectedBytes) + } +} + @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) @usableFromInline typealias TransactionBody = NIOThrowingAsyncSequenceProducer< diff --git a/Sources/AsyncHTTPClient/FoundationExtensions.swift b/Sources/AsyncHTTPClient/FoundationExtensions.swift index 545da756b..42b95b6d8 100644 --- a/Sources/AsyncHTTPClient/FoundationExtensions.swift +++ b/Sources/AsyncHTTPClient/FoundationExtensions.swift @@ -64,3 +64,14 @@ extension HTTPClient.Body { return self.bytes(data) } } + +@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) +extension HTTPClientResponse { + /// Response body as `Data`. + /// - Parameter maxBytes: The maximum number of bytes this method is allowed to accumulate. + /// - Returns: Bytes collected over time + public func data(upTo maxBytes: Int) async throws -> Data? { + var bytes = try await self.bytes(upTo: maxBytes) + return bytes.readData(length: bytes.readableBytes) + } +} diff --git a/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift b/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift index a30a8cf91..bfce896b6 100644 --- a/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift +++ b/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift @@ -835,6 +835,59 @@ final class AsyncAwaitEndToEndTests: XCTestCase { } } } + + func testResponseBytesHelper() { + XCTAsyncTest { + let bin = HTTPBin(.http2(compress: false)) { _ in HTTPEchoHandler() } + defer { XCTAssertNoThrow(try bin.shutdown()) } + let client = makeDefaultHTTPClient() + defer { XCTAssertNoThrow(try client.syncShutdown()) } + let logger = Logger(label: "HTTPClient", factory: StreamLogHandler.standardOutput(label:)) + var request = HTTPClientRequest(url: "https://localhost:\(bin.port)/") + request.method = .POST + request.body = .bytes(ByteBuffer(string: "1234")) + + guard let response = await XCTAssertNoThrowWithResult( + try await client.execute(request, deadline: .now() + .seconds(10), logger: logger) + ) else { return } + XCTAssertEqual(response.headers["content-length"], ["4"]) + guard let body = await XCTAssertNoThrowWithResult( + try await response.bytes(upTo: 3) + ) else { return } + XCTAssertEqual(body, ByteBuffer(string: "1234")) + + guard var responseNoContentLength = await XCTAssertNoThrowWithResult( + try await client.execute(request, deadline: .now() + .seconds(10), logger: logger) + ) else { return } + responseNoContentLength.headers.remove(name: "content-length") + guard let body2 = await XCTAssertNoThrowWithResult( + try await responseNoContentLength.bytes(upTo: 4) + ) else { return } + XCTAssertEqual(body2, ByteBuffer(string: "1234")) + } + } + + func testResponseBodyDataHelper() { + XCTAsyncTest { + let bin = HTTPBin(.http2(compress: false)) { _ in HTTPEchoHandler() } + defer { XCTAssertNoThrow(try bin.shutdown()) } + let client = makeDefaultHTTPClient() + defer { XCTAssertNoThrow(try client.syncShutdown()) } + let logger = Logger(label: "HTTPClient", factory: StreamLogHandler.standardOutput(label:)) + var request = HTTPClientRequest(url: "https://localhost:\(bin.port)/") + request.method = .POST + request.body = .bytes(ByteBuffer(string: "1234")) + + guard let response = await XCTAssertNoThrowWithResult( + try await client.execute(request, deadline: .now() + .seconds(10), logger: logger) + ) else { return } + XCTAssertEqual(response.headers["content-length"], ["4"]) + guard let bodyData = await XCTAssertNoThrowWithResult( + try await response.data(upTo: 4) + ) else { return } + XCTAssertEqual(bodyData, "1234".data(using: .utf8)) + } + } } struct AnySendableSequence: @unchecked Sendable {