more tests

Author: MahdiBM

Sources/PostgresNIO/New/PSQLRowStream.swift

@@ -296,14 +296,14 @@ final class PSQLRowStream: @unchecked Sendable {
         case .waitingForConsumer(let bufferState):
             switch bufferState {
             case .streaming(var buffer, let dataSource):
-                let promise = self.eventLoop.makePromise(of: Void.self)
+                let promise = self.eventLoop.makePromise(of: [PostgresRow].self)
 
                 buffer.removeAll()
-                self.downstreamState = .iteratingRows(onRow: { _ in }, promise, dataSource)
+                self.downstreamState = .waitingForAll([], promise, dataSource)
                 // immediately request more
                 dataSource.request(for: self)
 
-                return promise.futureResult
+                return promise.futureResult.map { _ in }
 
             case .finished(_, let summary):
                 self.downstreamState = .consumed(.success(summary))
@@ -316,13 +316,13 @@ final class PSQLRowStream: @unchecked Sendable {
         case .asyncSequence(let consumer, let dataSource, _):
             consumer.finish()
 
-            let promise = self.eventLoop.makePromise(of: Void.self)
+            let promise = self.eventLoop.makePromise(of: [PostgresRow].self)
 
-            self.downstreamState = .iteratingRows(onRow: { _ in }, promise, dataSource)
+            self.downstreamState = .waitingForAll([], promise, dataSource)
             // immediately request more
             dataSource.request(for: self)
 
-            return promise.futureResult
+            return promise.futureResult.map { _ in }
         case .consumed(.success):
             // already drained
             return self.eventLoop.makeSucceededVoidFuture()

Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift

@@ -10,7 +10,7 @@ final class PSQLRowStreamTests: XCTestCase {
     let logger = Logger(label: "PSQLRowStreamTests")
     let eventLoop = EmbeddedEventLoop()
 
-    func testEmptyStream() {
+    func testEmptyStreamAndDrainDoesNotThrowErrorAfterConsumption() {
         let stream = PSQLRowStream(
             source: .noRows(.success(.tag("INSERT 0 1"))),
             eventLoop: self.eventLoop,
@@ -20,10 +20,9 @@ final class PSQLRowStreamTests: XCTestCase {
         XCTAssertEqual(try stream.all().wait(), [])
         XCTAssertEqual(stream.commandTag, "INSERT 0 1")
 
-        // Test 'drain' works in this case
         XCTAssertNoThrow(try stream.drain().wait())
     }
-    
+
     func testFailedStream() {
         let stream = PSQLRowStream(
             source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))),
@@ -84,37 +83,37 @@ final class PSQLRowStreamTests: XCTestCase {
         )
         XCTAssertEqual(dataSource.hitDemand, 0)
         XCTAssertEqual(dataSource.hitCancel, 0)
-        
+
         stream.receive([
             [ByteBuffer(string: "0")],
             [ByteBuffer(string: "1")]
         ])
-        
+
         XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled")
-        
+
         // attach consumer
         let future = stream.all()
         XCTAssertEqual(dataSource.hitDemand, 1)
-        
+
         stream.receive([
             [ByteBuffer(string: "2")],
             [ByteBuffer(string: "3")]
         ])
         XCTAssertEqual(dataSource.hitDemand, 2)
-        
+
         stream.receive([
             [ByteBuffer(string: "4")],
             [ByteBuffer(string: "5")]
         ])
         XCTAssertEqual(dataSource.hitDemand, 3)
-        
+
         stream.receive(completion: .success("SELECT 2"))
-        
+
         var rows: [PostgresRow]?
         XCTAssertNoThrow(rows = try future.wait())
         XCTAssertEqual(rows?.count, 6)
     }
-    
+
     func testOnRowAfterStreamHasFinished() {
         let dataSource = CountingDataSource()
         let stream = PSQLRowStream(
@@ -240,6 +239,84 @@ final class PSQLRowStreamTests: XCTestCase {
         XCTAssertEqual(stream.commandTag, "SELECT 6")
     }
 
+    func testEmptyStreamDrainsSuccessfully() {
+        let stream = PSQLRowStream(
+            source: .noRows(.success(.tag("INSERT 0 1"))),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+
+        XCTAssertNoThrow(try stream.drain().wait())
+        XCTAssertEqual(stream.commandTag, "INSERT 0 1")
+    }
+
+    func testDrainAfterStreamHasFinished() {
+        let dataSource = CountingDataSource()
+        let stream = PSQLRowStream(
+            source: .stream(
+                [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)],
+                dataSource
+            ),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+        XCTAssertEqual(dataSource.hitDemand, 0)
+        XCTAssertEqual(dataSource.hitCancel, 0)
+
+        stream.receive([
+            [ByteBuffer(string: "0")],
+            [ByteBuffer(string: "1")]
+        ])
+
+        XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled")
+        stream.receive(completion: .success("SELECT 2"))
+
+        // attach consumer
+        XCTAssertNoThrow(try stream.drain().wait())
+        XCTAssertEqual(dataSource.hitDemand, 0) // TODO: Is this right?
+    }
+
+    func testDrainBeforeStreamHasFinished() {
+        let dataSource = CountingDataSource()
+        let stream = PSQLRowStream(
+            source: .stream(
+                [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)],
+                dataSource
+            ),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+        XCTAssertEqual(dataSource.hitDemand, 0)
+        XCTAssertEqual(dataSource.hitCancel, 0)
+
+        stream.receive([
+            [ByteBuffer(string: "0")],
+            [ByteBuffer(string: "1")]
+        ])
+
+        XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled")
+
+        // attach consumer
+        let future = stream.drain()
+        XCTAssertEqual(dataSource.hitDemand, 1)
+
+        stream.receive([
+            [ByteBuffer(string: "2")],
+            [ByteBuffer(string: "3")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 2)
+
+        stream.receive([
+            [ByteBuffer(string: "4")],
+            [ByteBuffer(string: "5")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 3)
+
+        stream.receive(completion: .success("SELECT 2"))
+
+        XCTAssertNoThrow(try future.wait())
+    }
+
     func makeColumnDescription(name: String, dataType: PostgresDataType, format: PostgresFormat) -> RowDescription.Column {
         RowDescription.Column(
             name: "test",