feat: improve task polling

src/posit/connect/tasks.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import time
-
 from typing_extensions import overload
 
 from . import resources
@@ -97,39 +95,21 @@ def update(self, *args, **kwargs) -> None:
         result = response.json()
         super().update(**result)
 
-    def wait_for(self, *, initial_wait: int = 1, max_wait: int = 10, backoff: float = 1.5) -> None:
+    def wait_for(self, *, wait: int = 1) -> None:
         """Wait for the task to finish.
 
         Parameters
         ----------
-        initial_wait : int, default 1
-            Initial wait time in seconds. First API request will use this as the wait parameter.
-        max_wait : int, default 10
-            Maximum wait time in seconds between polling requests.
-        backoff : float, default 1.5
-            Backoff multiplier for increasing wait times.
+        wait : int, default 1
+            Wait time in seconds between polling requests.
 
         Examples
         --------
         >>> task.wait_for()
         None
-
-        Notes
-        -----
-        This method implements an exponential backoff strategy to reduce the number of API calls
-        while waiting for long-running tasks. The first request uses the initial_wait value,
-        and subsequent requests increase the wait time by the backoff factor, up to max_wait. To disable exponential backoff, set backoff to 1.0.
         """
-        wait_time = initial_wait
-
         while not self.is_finished:
-            self.update()
-
-            # Wait client-side
-            time.sleep(wait_time)
-
-            # Calculate next wait time with backoff
-            wait_time = min(wait_time * backoff, max_wait)
+            self.update(wait=wait)
 
 
 class Tasks(resources.Resources):

tests/posit/connect/test_tasks.py

@@ -118,6 +118,7 @@ def test(self):
             responses.get(
                 f"https://connect.example/__api__/v1/tasks/{uid}",
                 json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": True},
+                match=[matchers.query_param_matcher({"wait": 1})],
             ),
         ]
 
@@ -127,35 +128,27 @@ def test(self):
         assert not task.is_finished
 
         # invoke
-        task.wait_for()
+        task.wait_for(wait=1)
 
         # assert
         assert task.is_finished
         assert mock_tasks_get[0].call_count == 1
         assert mock_tasks_get[1].call_count == 1
 
     @responses.activate
-    @mock.patch("time.sleep", autospec=True)
-    def test_exponential_backoff(self, mock_sleep):
+    def test_with_custom_wait(self):
         uid = "jXhOhdm5OOSkGhJw"
 
         # behavior
         mock_tasks_get = [
-            responses.get(
-                f"https://connect.example/__api__/v1/tasks/{uid}",
-                json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": False},
-            ),
-            responses.get(
-                f"https://connect.example/__api__/v1/tasks/{uid}",
-                json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": False},
-            ),
             responses.get(
                 f"https://connect.example/__api__/v1/tasks/{uid}",
                 json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": False},
             ),
             responses.get(
                 f"https://connect.example/__api__/v1/tasks/{uid}",
                 json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": True},
+                match=[matchers.query_param_matcher({"wait": 5})],
             ),
         ]
 
@@ -165,31 +158,19 @@ def test_exponential_backoff(self, mock_sleep):
         assert not task.is_finished
 
         # invoke
-        task.wait_for(initial_wait=1, max_wait=5, backoff=2.0)
+        task.wait_for(wait=5)
 
         # assert
         assert task.is_finished
         assert mock_tasks_get[0].call_count == 1
         assert mock_tasks_get[1].call_count == 1
 
-        # Verify sleep calls
-        mock_sleep.assert_has_calls([mock.call(1), mock.call(2), mock.call(4)], any_order=False)
-
     @responses.activate
-    @mock.patch("time.sleep", autospec=True)
-    def test_no_backoff(self, mock_sleep):
+    def test_immediate_completion(self):
         uid = "jXhOhdm5OOSkGhJw"
 
         # behavior
         mock_tasks_get = [
-            responses.get(
-                f"https://connect.example/__api__/v1/tasks/{uid}",
-                json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": False},
-            ),
-            responses.get(
-                f"https://connect.example/__api__/v1/tasks/{uid}",
-                json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": False},
-            ),
             responses.get(
                 f"https://connect.example/__api__/v1/tasks/{uid}",
                 json={**load_mock_dict(f"v1/tasks/{uid}.json"), "finished": True},
@@ -199,18 +180,14 @@ def test_no_backoff(self, mock_sleep):
         # setup
         c = connect.Client("https://connect.example", "12345")
         task = c.tasks.get(uid)
-        assert not task.is_finished
+        assert task.is_finished
 
         # invoke
-        task.wait_for(initial_wait=2, max_wait=5, backoff=1.0)
+        task.wait_for(wait=1)
 
         # assert
         assert task.is_finished
         assert mock_tasks_get[0].call_count == 1
-        assert mock_tasks_get[1].call_count == 1
-
-        # Verify sleep calls
-        mock_sleep.assert_has_calls([mock.call(2), mock.call(2)], any_order=False)
 
 
 class TestTasksGet: