diff --git a/bindings/java/src/test/java/org/openucx/jucx/UcpEndpointTest.java b/bindings/java/src/test/java/org/openucx/jucx/UcpEndpointTest.java
index 4fa5e1c315d..6768ab3884d 100644
--- a/bindings/java/src/test/java/org/openucx/jucx/UcpEndpointTest.java
+++ b/bindings/java/src/test/java/org/openucx/jucx/UcpEndpointTest.java
@@ -735,7 +735,8 @@ public void onSuccess(UcpRequest request) {
             sendData.getMemory().getAddress(), 2L, UcpConstants.UCP_AM_FLAG_PERSISTENT_DATA, null,
             new UcpRequestParams().setMemoryType(memType).setMemoryHandle(sendData.getMemory()));
 
-        while (!Arrays.stream(requests).allMatch(r -> (r != null) && r.isCompleted())) {
+        while (!Arrays.stream(requests).allMatch(r -> (r != null) && r.isCompleted()) ||
+             (persistantAmData.get() == null)) {
             worker1.progress();
             worker2.progress();
         }