added await_completed methods to base api

Author: sid-rl

src/runloop_api_client/resources/devboxes/disk_snapshots.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 import httpx
 
@@ -21,6 +21,9 @@
 from ...types.devboxes import disk_snapshot_list_params, disk_snapshot_update_params
 from ...types.devbox_snapshot_view import DevboxSnapshotView
 from ...types.devboxes.devbox_snapshot_async_status_view import DevboxSnapshotAsyncStatusView
+from ..._exceptions import RunloopError
+from ...lib.polling import PollingConfig, poll_until
+from ...lib.polling_async import async_poll_until
 
 __all__ = ["DiskSnapshotsResource", "AsyncDiskSnapshotsResource"]
 
@@ -239,6 +242,29 @@ def query_status(
             cast_to=DevboxSnapshotAsyncStatusView,
         )
 
+    def await_completed(
+        self,
+        id: str,
+        *,
+        polling_config: PollingConfig | None = None,
+        **request_options: Any,
+    ) -> DevboxSnapshotAsyncStatusView:
+        """Wait for a disk snapshot operation to complete."""
+
+        if not id:
+            raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
+
+        def is_terminal(result: DevboxSnapshotAsyncStatusView) -> bool:
+            return result.status in {"complete", "error"}
+
+        status = poll_until(lambda: self.query_status(id, **request_options), is_terminal, polling_config)
+
+        if status.status == "error":
+            message = status.error_message or "Unknown error"
+            raise RunloopError(f"Snapshot {id} failed: {message}")
+
+        return status
+
 
 class AsyncDiskSnapshotsResource(AsyncAPIResource):
     @cached_property
@@ -454,6 +480,33 @@ async def query_status(
             cast_to=DevboxSnapshotAsyncStatusView,
         )
 
+    async def await_completed(
+        self,
+        id: str,
+        *,
+        polling_config: PollingConfig | None = None,
+        **request_options: Any,
+    ) -> DevboxSnapshotAsyncStatusView:
+        """Wait asynchronously for a disk snapshot operation to complete."""
+
+        if not id:
+            raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
+
+        def is_terminal(result: DevboxSnapshotAsyncStatusView) -> bool:
+            return result.status in {"complete", "error"}
+
+        status = await async_poll_until(
+            lambda: self.query_status(id, **request_options),
+            is_terminal,
+            polling_config,
+        )
+
+        if status.status == "error":
+            message = status.error_message or "Unknown error"
+            raise RunloopError(f"Snapshot {id} failed: {message}")
+
+        return status
+
 
 class DiskSnapshotsResourceWithRawResponse:
     def __init__(self, disk_snapshots: DiskSnapshotsResource) -> None: