Two handy extensions to work with HTTPClientResponse body

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)
+    }
+}

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<Element>: @unchecked Sendable {