Add connection tests

Author: MahdiBM

Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift

@@ -217,6 +217,11 @@ extension PostgresFrontendMessage {
             preconditionFailure("TODO: Unimplemented")
         case .saslResponse:
             preconditionFailure("TODO: Unimplemented")
+        case .query:
+            guard let query = buffer.readNullTerminatedString() else {
+                throw PSQLPartialDecodingError.fieldNotDecodable(type: String.self)
+            }
+            return .query(.init(query: query))
         case .sync:
             return .sync
         case .terminate:

Tests/PostgresNIOTests/New/Extensions/PostgresFrontendMessage.swift

@@ -59,6 +59,11 @@ enum PostgresFrontendMessage: Equatable {
         }
     }
 
+    struct Query: Hashable {
+        /// The query string.
+        let query: String
+    }
+
     struct Parse: Hashable {
         /// The name of the destination prepared statement (an empty string selects the unnamed prepared statement).
         let preparedStatementName: String
@@ -179,6 +184,7 @@ enum PostgresFrontendMessage: Equatable {
     case saslInitialResponse(SASLInitialResponse)
     case saslResponse(SASLResponse)
     case sslRequest
+    case query(Query)
     case sync
     case startup(Startup)
     case terminate
@@ -194,6 +200,7 @@ enum PostgresFrontendMessage: Equatable {
         case password
         case saslInitialResponse
         case saslResponse
+        case query
         case sync
         case terminate
 
@@ -217,6 +224,8 @@ enum PostgresFrontendMessage: Equatable {
                 self = .saslInitialResponse
             case UInt8(ascii: "p"):
                 self = .saslResponse
+            case UInt8(ascii: "Q"):
+                self = .query
             case UInt8(ascii: "S"):
                 self = .sync
             case UInt8(ascii: "X"):
@@ -246,6 +255,8 @@ enum PostgresFrontendMessage: Equatable {
                 return UInt8(ascii: "p")
             case .saslResponse:
                 return UInt8(ascii: "p")
+            case .query:
+                return UInt8(ascii: "Q")
             case .sync:
                 return UInt8(ascii: "S")
             case .terminate:
@@ -283,6 +294,8 @@ extension PostgresFrontendMessage {
             preconditionFailure("SSL requests don't have an identifier")
         case .startup:
             preconditionFailure("Startup messages don't have an identifier")
+        case .query:
+            return .query
         case .sync:
             return .sync
         case .terminate:

Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift

@@ -31,7 +31,19 @@ class PSQLFrontendMessageTests: XCTestCase {
         XCTAssertEqual(PostgresFrontendMessage.ID.flush.rawValue, byteBuffer.readInteger(as: UInt8.self))
         XCTAssertEqual(4, byteBuffer.readInteger(as: Int32.self)) // payload length
     }
-    
+
+    func testEncodeQuery() {
+        var encoder = PostgresFrontendMessageEncoder(buffer: .init())
+        let query = "SELECT * FROM foo"
+        encoder.query(query)
+        var byteBuffer = encoder.flushBuffer()
+
+        XCTAssertEqual(byteBuffer.readableBytes, 23)
+        XCTAssertEqual(PostgresFrontendMessage.ID.query.rawValue, byteBuffer.readInteger(as: UInt8.self))
+        XCTAssertEqual(22, byteBuffer.readInteger(as: Int32.self)) // payload length
+        XCTAssertEqual([UInt8](query.utf8), byteBuffer.readBytes(length: 17))
+    }
+
     func testEncodeSync() {
         var encoder = PostgresFrontendMessageEncoder(buffer: .init())
         encoder.sync()

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -317,6 +317,40 @@ class PostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testCloseImmediatelyWithSimpleQuery() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+
+        try await withThrowingTaskGroup(of: Void.self) { [logger] taskGroup async throws -> () in
+            for _ in 1...2 {
+                taskGroup.addTask {
+                    try await connection.__simpleQuery("SELECT 1;", logger: logger)
+                }
+            }
+
+            let query = try await channel.waitForSimpleQueryRequest()
+            XCTAssertEqual(query.query, "SELECT 1;")
+
+            async let close: () = connection.close()
+
+            try await channel.closeFuture.get()
+            XCTAssertEqual(channel.isActive, false)
+
+            try await close
+
+            while let taskResult = await taskGroup.nextResult() {
+                switch taskResult {
+                case .success:
+                    XCTFail("Expected queries to fail")
+                case .failure(let failure):
+                    guard let error = failure as? PSQLError else {
+                        return XCTFail("Unexpected error type: \(failure)")
+                    }
+                    XCTAssertEqual(error.code, .clientClosedConnection)
+                }
+            }
+        }
+    }
+
     func testIfServerJustClosesTheErrorReflectsThat() async throws {
         let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
         let logger = self.logger
@@ -346,6 +380,35 @@ class PostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testIfServerJustClosesTheErrorReflectsThatInSimpleQuery() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+        let logger = self.logger
+
+        async let response = try await connection.__simpleQuery("SELECT 1;", logger: logger)
+
+        let query = try await channel.waitForSimpleQueryRequest()
+        XCTAssertEqual(query.query, "SELECT 1;")
+
+        try await channel.testingEventLoop.executeInContext { channel.pipeline.fireChannelInactive() }
+        try await channel.testingEventLoop.executeInContext { channel.pipeline.fireChannelUnregistered() }
+
+        do {
+            _ = try await response
+            XCTFail("Expected to throw")
+        } catch {
+            XCTAssertEqual((error as? PSQLError)?.code, .serverClosedConnection)
+        }
+
+        // retry on same connection
+
+        do {
+            _ = try await connection.__simpleQuery("SELECT 1;", logger: self.logger)
+            XCTFail("Expected to throw")
+        } catch {
+            XCTAssertEqual((error as? PSQLError)?.code, .serverClosedConnection)
+        }
+    }
+
     struct TestPrepareStatement: PostgresPreparedStatement {
         static let sql = "SELECT datname FROM pg_stat_activity WHERE state = $1"
         typealias Row = String
@@ -692,6 +755,14 @@ extension NIOAsyncTestingChannel {
         return UnpreparedRequest(parse: parse, describe: describe, bind: bind, execute: execute)
     }
 
+    func waitForSimpleQueryRequest() async throws -> PostgresFrontendMessage.Query {
+        let query = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self)
+        guard case .query(let query) = query else {
+            fatalError()
+        }
+        return query
+    }
+
     func waitForPrepareRequest() async throws -> PrepareRequest {
         let parse = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self)
         let describe = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self)