feat: add ActionsClient.wait_for function

This function allows the users to wait for multiple actions in an efficient way. All actions are queried using a single call, which reduce the potential for running into rate limits.

Author: jooola

hcloud/_utils.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from itertools import islice
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def batched(iterable: Iterable[T], size: int) -> Iterator[tuple[T, ...]]:
+    """
+    Returns a batch of the provided size from the provided iterable.
+    """
+    iterator = iter(iterable)
+    while True:
+        batch = tuple(islice(iterator, size))
+        if not batch:
+            break
+        yield batch

hcloud/actions/client.py

@@ -2,8 +2,9 @@
 
 import time
 import warnings
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any, Callable, NamedTuple
 
+from .._utils import batched
 from ..core import BoundModelBase, ClientEntityBase, Meta
 from .domain import Action, ActionFailedException, ActionTimeoutException
 
@@ -129,6 +130,100 @@ class ActionsClient(ResourceActionsClient):
     def __init__(self, client: Client):
         super().__init__(client, None)
 
+    def _get_list_by_ids(self, ids: list[int]) -> list[BoundAction]:
+        """
+        Get a list of Actions by their IDs.
+
+        :param ids: List of Action IDs to get.
+        :raises ValueError: Raise when Action IDs were not found.
+        :return: List of Actions.
+        """
+        actions: list[BoundAction] = []
+
+        for ids_batch in batched(ids, 25):
+            params: dict[str, Any] = {
+                "id": ids_batch,
+            }
+
+            response = self._client.request(
+                method="GET",
+                url="/actions",
+                params=params,
+            )
+
+            actions.extend(
+                BoundAction(self._client.actions, action_data)
+                for action_data in response["actions"]
+            )
+
+        if len(ids) != len(actions):
+            found_ids = [a.id for a in actions]
+            not_found_ids = list(set(ids) - set(found_ids))
+
+            raise ValueError(f"actions not found: {', '.join(not_found_ids)}")
+
+        return actions
+
+    def wait_for_function(
+        self,
+        handle_update: Callable[[BoundAction], None],
+        actions: list[Action | BoundAction],
+    ) -> list[BoundAction]:
+        """
+        Waits until all Actions succeed by polling the API at the interval defined by
+        the client's poll interval and function. An Action is considered as complete
+        when its status is either "success" or "error".
+
+        The handle_update callback is called every time an Action is updated.
+
+        :param handle_update: Function called every time an Action is updated.
+        :param actions: List of Actions to wait for.
+        :raises: ActionFailedException when an Action failed.
+        :return: List of succeeded Actions.
+        """
+        running_ids = [a.id for a in actions]
+
+        completed: list[BoundAction] = []
+
+        retries = 0
+        while len(running_ids):
+            # pylint: disable=protected-access
+            time.sleep(self._client._poll_interval_func(retries))
+            retries += 1
+
+            updates = self._get_list_by_ids(running_ids)
+
+            for update in updates:
+                if update.status != Action.STATUS_RUNNING:
+                    running_ids.remove(update.id)
+                    completed.append(update)
+
+                handle_update(update)
+
+        return completed
+
+    def wait_for(
+        self,
+        actions: list[Action | BoundAction],
+    ) -> list[BoundAction]:
+        """
+        Waits until all Actions succeed by polling the API at the interval defined by
+        the client's poll interval and function. An Action is considered as complete
+        when its status is either "success" or "error".
+
+        If a single Action fails, the function will stop waiting and raise ActionFailedException.
+
+        :param actions: List of Actions to wait for.
+        :raises: ActionFailedException when an Action failed.
+        :return: List of succeeded Actions.
+        """
+
+        def handle_update(update: BoundAction):
+            if update.status == Action.STATUS_ERROR:
+                raise ActionFailedException(action=update)
+
+        return self.wait_for_function(handle_update, actions)
+
     def get_list(
         self,
         status: list[str] | None = None,

tests/unit/actions/test_client.py

@@ -197,3 +197,64 @@ def test_get_all(self, actions_client, generic_action_list, params):
         assert action2._client == actions_client._client.actions
         assert action2.id == 2
         assert action2.command == "stop_server"
+
+    def test_wait_for(self, actions_client: ActionsClient):
+        actions = [Action(id=1), Action(id=2)]
+
+        # Speed up test by not really waiting
+        actions_client._client._poll_interval_func = mock.MagicMock()
+        actions_client._client._poll_interval_func.return_value = 0.1
+
+        actions_client._client.request.side_effect = [
+            {
+                "actions": [
+                    {"id": 1, "status": "running"},
+                    {"id": 2, "status": "success"},
+                ]
+            },
+            {
+                "actions": [
+                    {"id": 1, "status": "success"},
+                ]
+            },
+        ]
+
+        actions = actions_client.wait_for(actions)
+
+        actions_client._client.request.assert_has_calls(
+            [
+                mock.call(method="GET", url="/actions", params={"id": (1, 2)}),
+                mock.call(method="GET", url="/actions", params={"id": (1,)}),
+            ]
+        )
+
+        assert len(actions) == 2
+
+    def test_wait_for_error(self, actions_client: ActionsClient):
+        actions = [Action(id=1), Action(id=2)]
+
+        # Speed up test by not really waiting
+        actions_client._client._poll_interval_func = mock.MagicMock()
+        actions_client._client._poll_interval_func.return_value = 0.1
+
+        actions_client._client.request.side_effect = [
+            {
+                "actions": [
+                    {"id": 1, "status": "running"},
+                    {
+                        "id": 2,
+                        "status": "error",
+                        "error": {"code": "failed", "message": "Action failed"},
+                    },
+                ]
+            },
+        ]
+
+        with pytest.raises(ActionFailedException):
+            actions_client.wait_for(actions)
+
+        actions_client._client.request.assert_has_calls(
+            [
+                mock.call(method="GET", url="/actions", params={"id": (1, 2)}),
+            ]
+        )

tests/unit/test_utils.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from hcloud._utils import batched
+
+
+def test_batched():
+    assert list(o for o in batched([1, 2, 3, 4, 5], 2)) == [(1, 2), (3, 4), (5,)]