diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift index f98e64bc..16441dfc 100644 --- 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))