diff --git a/client/rest-high-level/src/test/java/org/elasticsearch/client/ESRestHighLevelClientTestCase.java b/client/rest-high-level/src/test/java/org/elasticsearch/client/ESRestHighLevelClientTestCase.java
index 83d7f19cd811f..d9ead045bcf97 100644
--- a/client/rest-high-level/src/test/java/org/elasticsearch/client/ESRestHighLevelClientTestCase.java
+++ b/client/rest-high-level/src/test/java/org/elasticsearch/client/ESRestHighLevelClientTestCase.java
@@ -351,7 +351,8 @@ protected static TaskId findTaskToRethrottle(String actionName, String descripti
             TaskGroup taskGroup = taskGroups.get(0);
             assertThat(taskGroup.getChildTasks(), empty());
             // check that the task initialized enough that it can rethrottle too.
-            if (((RawTaskStatus) taskGroup.getTaskInfo().getStatus()).toMap().containsKey("batches")) {
+            Map<String, Object> statusMap = ((RawTaskStatus) taskGroup.getTaskInfo().getStatus()).toMap();
+            if (statusMap.get("batches").equals(1)) {
                 return taskGroup.getTaskInfo().getTaskId();
             }
         } while (System.nanoTime() - start < TimeUnit.SECONDS.toNanos(10));