comments and tests

hummingbird-project · Dec 13, 2024 · e4f3305 · e4f3305
1 parent d4d2b76
commit e4f3305
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/Sources/HummingbirdCore/Request/RequestBodyMergedWithUnderlyingRequestPartIterator.swift b/Sources/HummingbirdCore/Request/RequestBodyMergedWithUnderlyingRequestPartIterator.swift
@@ -15,6 +15,11 @@
 import NIOCore
 import NIOHTTPTypes
 
+/// AsyncSequence used by consumeWithInboundCloseHandler
+///
+/// It will provide the buffers output by the ResponseBody and when that finishes will start
+/// iterating what is left of the underlying request part stream, and continue iterating until
+/// it hits the next head
 struct RequestBodyMergedWithUnderlyingRequestPartIterator<Base: AsyncSequence>: AsyncSequence where Base.Element == ByteBuffer {
     typealias Element = HTTPRequestPart
     let base: Base
@@ -24,6 +29,7 @@ struct RequestBodyMergedWithUnderlyingRequestPartIterator<Base: AsyncSequence>:
         enum CurrentAsyncIterator {
             case base(Base.AsyncIterator, underlying: NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator)
             case underlying(NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator)
+            case done
         }
         var current: CurrentAsyncIterator
 
@@ -41,10 +47,20 @@ struct RequestBodyMergedWithUnderlyingRequestPartIterator<Base: AsyncSequence>:
                     self.current = .underlying(underlying)
                     return .end(nil)
                 }
+
             case .underlying(var underlying):
-                let element = try await underlying.next()
+                while true {
+                    let part = try await underlying.next()
+                    if case .head = part {
+                        self.current = .done
+                        return part
+                    }
+                }
                 self.current = .underlying(underlying)
-                return element
+                return nil
+
+            case .done:
+                return nil
             }
         }
     }

diff --git a/Tests/HummingbirdTests/ApplicationTests.swift b/Tests/HummingbirdTests/ApplicationTests.swift
@@ -894,7 +894,7 @@ final class ApplicationTests: XCTestCase {
         }
     }
 
-    /// Test consumeWithInboundHandler
+    /// Test consumeWithInboundHandler after having collected the Request body
     @available(macOS 15, iOS 18, tvOS 18, *)
     func testConsumeWithInboundHandlerAfterCollect() async throws {
         let router = Router()
@@ -924,15 +924,15 @@ final class ApplicationTests: XCTestCase {
         }
     }
 
-    /// Test consumeWithInboundHandler
+    /// Test consumeWithInboundHandler after having replaced Request.body with a new streamed RequestBody
     @available(macOS 15, iOS 18, tvOS 18, *)
     func testConsumeWithInboundHandlerAfterReplacingBody() async throws {
         let router = Router()
         router.post("streaming") { request, context -> Response in
             var request = request
             request.body = .init(
                 asyncSequence: request.body.map {
-                    var view = $0.readableBytesView.map { $0 ^ 255 }
+                    let view = $0.readableBytesView.map { $0 ^ 255 }
                     return ByteBuffer(bytes: view)
                 }
             )