one more test

Author: MahdiBM

Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift

@@ -36,11 +36,12 @@ final class PSQLRowStreamTests: XCTestCase {
             XCTAssertEqual($0 as? PSQLError, expectedError)
         }
 
+        // Drain should work
         XCTAssertThrowsError(try stream.drain().wait()) {
             XCTAssertEqual($0 as? PSQLError, expectedError)
         }
     }
-    
+
     func testGetArrayAfterStreamHasFinished() {
         let dataSource = CountingDataSource()
         let stream = PSQLRowStream(
@@ -250,6 +251,20 @@ final class PSQLRowStreamTests: XCTestCase {
         XCTAssertEqual(stream.commandTag, "INSERT 0 1")
     }
 
+    func testDrainFailedStream() {
+        let stream = PSQLRowStream(
+            source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+
+        let expectedError = PSQLError.serverClosedConnection(underlying: nil)
+
+        XCTAssertThrowsError(try stream.drain().wait()) {
+            XCTAssertEqual($0 as? PSQLError, expectedError)
+        }
+    }
+
     func testDrainAfterStreamHasFinished() {
         let dataSource = CountingDataSource()
         let stream = PSQLRowStream(