more tests

Author: MahdiBM

Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift

@@ -19,6 +19,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() {
@@ -27,9 +30,15 @@ final class PSQLRowStreamTests: XCTestCase {
             eventLoop: self.eventLoop,
             logger: self.logger
         )
-        
+
+        let expectedError = PSQLError.serverClosedConnection(underlying: nil)
+
         XCTAssertThrowsError(try stream.all().wait()) {
-            XCTAssertEqual($0 as? PSQLError, .serverClosedConnection(underlying: nil))
+            XCTAssertEqual($0 as? PSQLError, expectedError)
+        }
+
+        XCTAssertThrowsError(try stream.drain().wait()) {
+            XCTAssertEqual($0 as? PSQLError, expectedError)
         }
     }