diff --git a/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse+data.swift b/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse+data.swift new file mode 100644 index 000000000..44099a420 --- /dev/null +++ b/Sources/AsyncHTTPClient/AsyncAwait/HTTPClientResponse+data.swift @@ -0,0 +1,41 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the AsyncHTTPClient open source project +// +// Copyright (c) 2021 Apple Inc. and the AsyncHTTPClient project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +import NIOCore +import NIOHTTP1 +import Foundation +internal import NIOFoundationCompat + +@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: min(expectedBytes, maxBytes)) + } +} + + +@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 56a08b852..9f3da4085 100644 --- a/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift +++ b/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift @@ -1019,6 +1019,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: 5) + ) 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 {