TestKit: allow new transaction after tx failure

robsdedude · Jan 20, 2024 · cb56fcb · cb56fcb
1 parent a35f401
commit cb56fcb
Showing 1 changed file with 17 additions and 3 deletions.
diff --git a/testkit_backend/src/testkit_backend/session_holder.rs b/testkit_backend/src/testkit_backend/session_holder.rs
@@ -631,9 +631,19 @@ impl SessionHolderRunner {
                                     }
 
                                     command @ (Command::BeginTransaction(_)
-                                    | Command::TransactionFunction(_)|Command::AutoCommit(_)) => command.reply_error(
-                                        tx_res, session_already_executing_tx_error()
-                                    ),
+                                    | Command::TransactionFunction(_)|Command::AutoCommit(_)) => {
+                                        if known_transactions
+                                            .get(&id)
+                                            .map(|tx| matches!(tx, TxFailState::Failed))
+                                            .unwrap_or(false) {
+                                            // transaction failed, therefore, we allow to start a new one
+                                            _ = buffered_command.insert(command);
+                                            return Ok(());
+                                        }
+                                        command.reply_error(
+                                            tx_res, session_already_executing_tx_error(),
+                                        )
+                                    }
                                     command @ (Command::RetryablePositive(_)
                                     | Command::RetryableNegative(_)) => {
                                         command.reply_error(
@@ -691,6 +701,10 @@ impl SessionHolderRunner {
                                     .unwrap();
                                 return;
                             }
+                            command @ (Command::BeginTransaction(_)
+                                | Command::TransactionFunction(_)|Command::AutoCommit(_)) => {
+                                _ = buffered_command.insert(command);
+                            }
                             _ => panic!(
                                 "left transaction function with unexpected buffered command {command:?}"
                             ),