Return the finished Task from waitForTask and add a Client timeout overload

src/main/java/com/meilisearch/sdk/Client.java

@@ -356,10 +356,29 @@ public TaskInfo deleteTasks(DeleteTasksQuery param) throws MeilisearchException
      * Waits for a task to be processed
      *
      * @param uid Identifier of the requested Task
+     * @return Task in its final state (succeeded, failed or canceled)
      * @throws MeilisearchException if an error occurs or if timeout is reached
+     * @see <a href="https://www.meilisearch.com/docs/reference/api/tasks#task-status">API
+     *     specification</a>
      */
-    public void waitForTask(int uid) throws MeilisearchException {
-        this.tasksHandler.waitForTask(uid);
+    public Task waitForTask(int uid) throws MeilisearchException {
+        return this.tasksHandler.waitForTask(uid);
+    }
+
+    /**
+     * Waits for a task to be processed
+     *
+     * @param uid Identifier of the requested Task
+     * @param timeoutInMs number of milliseconds before throwing an Exception
+     * @param intervalInMs number of milliseconds before requesting the status again
+     * @return Task in its final state (succeeded, failed or canceled)
+     * @throws MeilisearchException if an error occurs or if timeout is reached
+     * @see <a href="https://www.meilisearch.com/docs/reference/api/tasks#task-status">API
+     *     specification</a>
+     */
+    public Task waitForTask(int uid, int timeoutInMs, int intervalInMs)
+            throws MeilisearchException {
+        return this.tasksHandler.waitForTask(uid, timeoutInMs, intervalInMs);
     }
 
     /**

src/main/java/com/meilisearch/sdk/Index.java

@@ -1307,12 +1307,13 @@ public TasksResults getTasks(TasksQuery param) throws MeilisearchException {
      * Waits for a task to be processed
      *
      * @param taskId Identifier of the requested Task
+     * @return Task in its final state (succeeded, failed or canceled)
      * @throws MeilisearchException if an error occurs or if timeout is reached
      * @see <a href="https://www.meilisearch.com/docs/reference/api/tasks#task-status">API
      *     specification</a>
      */
-    public void waitForTask(int taskId) throws MeilisearchException {
-        this.tasksHandler.waitForTask(taskId, 5000, 50);
+    public Task waitForTask(int taskId) throws MeilisearchException {
+        return this.tasksHandler.waitForTask(taskId);
     }
 
     /**
@@ -1321,13 +1322,14 @@ public void waitForTask(int taskId) throws MeilisearchException {
      * @param taskId ID of the index update
      * @param timeoutInMs number of milliseconds before throwing an Exception
      * @param intervalInMs number of milliseconds before requesting the status again
+     * @return Task in its final state (succeeded, failed or canceled)
      * @throws MeilisearchException if an error occurs or if timeout is reached
      * @see <a href="https://www.meilisearch.com/docs/reference/api/tasks#task-status">API
      *     specification</a>
      */
-    public void waitForTask(int taskId, int timeoutInMs, int intervalInMs)
+    public Task waitForTask(int taskId, int timeoutInMs, int intervalInMs)
             throws MeilisearchException {
-        this.tasksHandler.waitForTask(taskId, timeoutInMs, intervalInMs);
+        return this.tasksHandler.waitForTask(taskId, timeoutInMs, intervalInMs);
     }
 
     /**

src/main/java/com/meilisearch/sdk/TasksHandler.java

@@ -6,14 +6,16 @@
 import com.meilisearch.sdk.model.*;
 import com.meilisearch.sdk.model.batch.req.BatchesQuery;
 import com.meilisearch.sdk.model.batch.res.Batch;
-import java.util.Date;
 
 /**
  * Class covering the Meilisearch Task API
  *
  * @see <a href="https://www.meilisearch.com/docs/reference/api/tasks">API specification</a>
  */
 public class TasksHandler {
+    static final int DEFAULT_WAIT_TIMEOUT_MS = 5000;
+    static final int DEFAULT_WAIT_INTERVAL_MS = 50;
+
     private final HttpClient httpClient;
 
     /**
@@ -124,10 +126,11 @@ TaskInfo deleteTasks(DeleteTasksQuery param) throws MeilisearchException {
      * Waits for a task to be processed
      *
      * @param taskUid Identifier of the Task
+     * @return Task in its final state (succeeded, failed or canceled)
      * @throws MeilisearchException if timeout is reached
      */
-    void waitForTask(int taskUid) throws MeilisearchException {
-        this.waitForTask(taskUid, 5000, 50);
+    Task waitForTask(int taskUid) throws MeilisearchException {
+        return this.waitForTask(taskUid, DEFAULT_WAIT_TIMEOUT_MS, DEFAULT_WAIT_INTERVAL_MS);
     }
 
     /**
@@ -136,28 +139,35 @@ void waitForTask(int taskUid) throws MeilisearchException {
      * @param taskUid Identifier of the Task
      * @param timeoutInMs number of milliseconds before throwing an Exception
      * @param intervalInMs number of milliseconds before requesting the status again
+     * @return Task in its final state (succeeded, failed or canceled)
      * @throws MeilisearchException if timeout is reached
      */
-    void waitForTask(int taskUid, int timeoutInMs, int intervalInMs) throws MeilisearchException {
-        Task task;
-        TaskStatus status = null;
-        long startTime = new Date().getTime();
-        long elapsedTime = 0;
-
-        while (status == null
-                || (status.equals(TaskStatus.ENQUEUED) || status.equals(TaskStatus.PROCESSING))) {
-            if (elapsedTime >= timeoutInMs) {
-                throw new MeilisearchTimeoutException();
+    Task waitForTask(int taskUid, int timeoutInMs, int intervalInMs) throws MeilisearchException {
+        long deadline = System.currentTimeMillis() + timeoutInMs;
+        while (true) {
+            Task task = this.getTask(taskUid);
+            TaskStatus status = task.getStatus();
+            if (status != TaskStatus.ENQUEUED && status != TaskStatus.PROCESSING) {
+                return task;
+            }
+            long remainingMs = deadline - System.currentTimeMillis();
+            if (remainingMs <= 0) {
+                throw new MeilisearchTimeoutException(
+                        "Task "
+                                + taskUid
+                                + " not finished after "
+                                + timeoutInMs
+                                + "ms (last status: "
+                                + status
+                                + ")");
             }
-            task = this.getTask(taskUid);
-            status = task.getStatus();
             try {
-                Thread.sleep(intervalInMs);
+                // never sleep past the deadline, even when intervalInMs exceeds timeoutInMs
+                Thread.sleep(Math.min(intervalInMs, remainingMs));
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
-                throw new MeilisearchTimeoutException();
+                throw new MeilisearchTimeoutException(e);
             }
-            elapsedTime = new Date().getTime() - startTime;
         }
     }
 

src/test/java/com/meilisearch/integration/TasksTest.java

@@ -2,10 +2,12 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.blankOrNullString;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
@@ -14,6 +16,7 @@
 import com.meilisearch.integration.classes.AbstractIT;
 import com.meilisearch.integration.classes.TestData;
 import com.meilisearch.sdk.Index;
+import com.meilisearch.sdk.exceptions.MeilisearchTimeoutException;
 import com.meilisearch.sdk.model.*;
 import com.meilisearch.sdk.utils.Movie;
 import java.time.Instant;
@@ -340,9 +343,77 @@ public void testWaitForTaskTimoutInMs() throws Exception {
         Index index = client.index(indexUid);
 
         TaskInfo task = index.addDocuments(this.testData.getRaw());
+
+        MeilisearchTimeoutException e =
+                assertThrows(
+                        MeilisearchTimeoutException.class,
+                        () -> index.waitForTask(task.getTaskUid(), 0, 50));
+        assertThat(e.getMessage(), containsString("Task " + task.getTaskUid()));
+        assertThat(e.getMessage(), containsString("0ms"));
+
+        index.waitForTask(task.getTaskUid());
+    }
+
+    /** Test waitForTask does not sleep past the timeout when intervalInMs exceeds it */
+    @Test
+    public void testWaitForTaskIntervalLongerThanTimeout() throws Exception {
+        String indexUid = "WaitForTaskIntervalLongerThanTimeout";
+        Index index = client.index(indexUid);
+        TaskInfo task = index.addDocuments(this.testData.getRaw());
+
+        long start = System.currentTimeMillis();
+        try {
+            index.waitForTask(task.getTaskUid(), 100, 10000);
+        } catch (MeilisearchTimeoutException ignored) {
+            // either outcome is fine, only the elapsed time matters
+        }
+
+        assertThat(System.currentTimeMillis() - start, is(lessThan(5000L)));
+
         index.waitForTask(task.getTaskUid());
+    }
+
+    /** Test waitForTask returns the finished task */
+    @Test
+    public void testWaitForTaskReturnsTask() throws Exception {
+        String indexUid = "WaitForTaskReturnsTask";
+        TaskInfo response = client.createIndex(indexUid);
+
+        Task task = client.waitForTask(response.getTaskUid());
+
+        assertThat(task.getUid(), is(equalTo(response.getTaskUid())));
+        assertThat(task.getStatus(), is(equalTo(TaskStatus.SUCCEEDED)));
+        assertThat(task.getFinishedAt(), is(notNullValue()));
+
+        client.deleteIndex(indexUid);
+    }
+
+    /** Test waitForTask returns a failed task instead of throwing */
+    @Test
+    public void testWaitForTaskReturnsFailedTask() throws Exception {
+        String indexUid = "WaitForTaskReturnsFailedTask";
+        client.waitForTask(client.createIndex(indexUid).getTaskUid());
+
+        TaskInfo response = client.createIndex(indexUid);
+        Task task = client.waitForTask(response.getTaskUid());
+
+        assertThat(task.getStatus(), is(equalTo(TaskStatus.FAILED)));
+        assertThat(task.getError().getCode(), is(equalTo("index_already_exists")));
+
+        client.deleteIndex(indexUid);
+    }
+
+    /** Test Client.waitForTask with timeoutInMs and intervalInMs */
+    @Test
+    public void testClientWaitForTaskTimeoutInMs() throws Exception {
+        String indexUid = "ClientWaitForTaskTimeoutInMs";
+        TaskInfo response = client.createIndex(indexUid);
+
+        Task task = client.waitForTask(response.getTaskUid(), 10000, 50);
+
+        assertThat(task.getStatus(), is(equalTo(TaskStatus.SUCCEEDED)));
 
-        assertThrows(Exception.class, () -> index.waitForTask(task.getTaskUid(), 0, 50));
+        client.deleteIndex(indexUid);
     }
 
     /** Test Tasks with Jackson Json Handler */