Add another cancellation error

vapor · Aug 26, 2024 · 26fe020 · 26fe020
1 parent 807d8a0
commit 26fe020
Showing 1 changed file with 43 additions and 11 deletions.
diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift
@@ -57,14 +57,10 @@ class SimpleQueryStateMachineTests: XCTestCase {
         let input: [RowDescription.Column] = [
             .init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
         ]
-        let expected: [RowDescription.Column] = input.map {
-            .init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
-                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .text)
-        }
 
         XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
         let row1: DataRow = [ByteBuffer(string: "test1")]
-        let result = QueryResult(value: .rowDescription(expected), logger: queryContext.logger)
+        let result = QueryResult(value: .rowDescription(input), logger: queryContext.logger)
         XCTAssertEqual(state.dataRowReceived(row1), .succeedQuery(promise, with: result))
         XCTAssertEqual(state.channelReadComplete(), .forwardRows([row1]))
         XCTAssertEqual(state.readEventCaught(), .wait)
@@ -128,7 +124,7 @@ class SimpleQueryStateMachineTests: XCTestCase {
                        .failQuery(promise, with: psqlError, cleanupContext: .init(action: .close, tasks: [], error: psqlError, closePromise: nil)))
     }
 
-    func testQueryIsCancelledImmediatly() {
+    func testQueryIsCancelledImmediately() {
         var state = ConnectionStateMachine.readyForQuery()
 
         let logger = Logger.psqlTest
@@ -183,6 +179,46 @@ class SimpleQueryStateMachineTests: XCTestCase {
         XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
     }
 
+    func testQueryIsCancelledWithReadPendingWhileStreaming() {
+        var state = ConnectionStateMachine.readyForQuery()
+
+        let logger = Logger.psqlTest
+        let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
+        promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
+        let query = "SELECT version()"
+        let queryContext = SimpleQueryContext(query: query, logger: logger, promise: promise)
+
+        XCTAssertEqual(state.enqueue(task: .simpleQuery(queryContext)), .sendQuery(query))
+
+        // We need to ensure that even though the row description from the wire says that we
+        // will receive data in `.text` format, we will actually receive it in binary format,
+        // since we requested it in binary with our bind message.
+        let input: [RowDescription.Column] = [
+            .init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
+        ]
+
+        XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
+        let row1: DataRow = [ByteBuffer(string: "test1")]
+        let result = QueryResult(value: .rowDescription(input), logger: queryContext.logger)
+        XCTAssertEqual(state.dataRowReceived(row1), .succeedQuery(promise, with: result))
+        XCTAssertEqual(state.cancelQueryStream(), .forwardStreamError(.queryCancelled, read: false, cleanupContext: nil))
+        XCTAssertEqual(state.channelReadComplete(), .wait)
+        XCTAssertEqual(state.readEventCaught(), .read)
+
+        XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test2")]), .wait)
+        XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test3")]), .wait)
+        XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test4")]), .wait)
+        XCTAssertEqual(state.channelReadComplete(), .wait)
+        XCTAssertEqual(state.readEventCaught(), .read)
+
+        XCTAssertEqual(state.channelReadComplete(), .wait)
+        XCTAssertEqual(state.readEventCaught(), .read)
+
+        XCTAssertEqual(state.commandCompletedReceived("SELECT 2"), .wait)
+        XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
+    }
+
+
     func testCancelQueryAfterServerError() {
         var state = ConnectionStateMachine.readyForQuery()
 
@@ -200,13 +236,9 @@ class SimpleQueryStateMachineTests: XCTestCase {
         let input: [RowDescription.Column] = [
             .init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
         ]
-        let expected: [RowDescription.Column] = input.map {
-            .init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
-                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .text)
-        }
 
         XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
-        let result = QueryResult(value: .rowDescription(expected), logger: queryContext.logger)
+        let result = QueryResult(value: .rowDescription(input), logger: queryContext.logger)
         let row1: DataRow = [ByteBuffer(string: "test1")]
         XCTAssertEqual(state.dataRowReceived(row1), .succeedQuery(promise, with: result))