From a28236b1782ca46484c0af16ef6b6eb36731fbc2 Mon Sep 17 00:00:00 2001
From: rbourbonnavarro <rodrigobourbon44@gmail.com>
Date: Fri, 20 Dec 2024 16:56:14 -0300
Subject: [PATCH] W-17430963: fix flaky test

---
 .../TransactionRolledBackByOwnerTestCase.java            | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/integration/src/test/java/org/mule/test/integration/transaction/TransactionRolledBackByOwnerTestCase.java b/integration/src/test/java/org/mule/test/integration/transaction/TransactionRolledBackByOwnerTestCase.java
index 560a74f933..f7a615a12f 100644
--- a/integration/src/test/java/org/mule/test/integration/transaction/TransactionRolledBackByOwnerTestCase.java
+++ b/integration/src/test/java/org/mule/test/integration/transaction/TransactionRolledBackByOwnerTestCase.java
@@ -8,6 +8,7 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.iterableWithSize;
 import static org.junit.Assert.fail;
@@ -337,7 +338,7 @@ public void checkRollback() {
         }
       }
 
-      assertStatesArrived();
+      assertStatesArrived(flowExecution.expectedFinalState);
       assertCorrectStates(flowExecution);
       if (flowExecution.globalHandlerExecutionsBeforeRollback != null) {
         assertTransactionRolledBackAtTheRightHandler(flowExecution);
@@ -345,19 +346,19 @@ public void checkRollback() {
     });
   }
 
-  private void assertStatesArrived() {
+  private void assertStatesArrived(String expectedFinalState) {
     PollingProber pollingProber = new PollingProber(RECEIVE_TIMEOUT, POLL_DELAY_MILLIS);
     pollingProber.check(new JUnitProbe() {
 
       @Override
       protected boolean test() {
-        assertThat(states.size(), greaterThanOrEqualTo(2));
+        assertThat(states, hasItem(expectedFinalState));
         return true;
       }
 
       @Override
       public String describeFailure() {
-        return "States did not arrive";
+        return "Final state " + expectedFinalState + " not received " + states;
       }
     });
   }