From 6357d9f5a2521016584b579de041e04686b40aa4 Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 17 Aug 2026 23:48:42 +0000 Subject: [PATCH] fix(devboxes): clean up create timeout --- .../resources/devboxes/devboxes.py | 40 ++++++--- src/runloop_api_client/sdk/_types.py | 9 +- tests/api_resources/test_devboxes.py | 90 ++++++++++++++++++- tests/sdk/test_async_ops.py | 2 + tests/sdk/test_ops.py | 2 + 5 files changed, 129 insertions(+), 14 deletions(-) diff --git a/src/runloop_api_client/resources/devboxes/devboxes.py b/src/runloop_api_client/resources/devboxes/devboxes.py index 65d0dcbd3..4e51f7a95 100644 --- a/src/runloop_api_client/resources/devboxes/devboxes.py +++ b/src/runloop_api_client/resources/devboxes/devboxes.py @@ -77,7 +77,7 @@ AsyncDiskSnapshotsCursorIDPage, ) from ..._exceptions import RunloopError, APIStatusError, APIConnectionError -from ...lib.polling import PollingConfig, poll_until +from ...lib.polling import PollingConfig, PollingTimeout, poll_until from ..._base_client import AsyncPaginator, make_request_options from .disk_snapshots import ( DiskSnapshotsResource, @@ -471,6 +471,7 @@ def create_and_await_running( mounts: Optional[Iterable[Mount]] | Omit = omit, name: Optional[str] | Omit = omit, polling_config: PollingConfig | None = None, + shutdown_on_timeout: bool = True, secrets: Optional[Dict[str, str]] | Omit = omit, snapshot_id: Optional[str] | Omit = omit, tunnel: Optional[devbox_create_params.Tunnel] | Omit = omit, @@ -489,9 +490,11 @@ def create_and_await_running( Args: create_args: Arguments to pass to the `create` method. See the `create` method for detailed documentation. request_args: Optional request arguments including polling configuration and additional request options + shutdown_on_timeout: Shutdown the created devbox if waiting for running state times out. Returns: - The devbox in running state + The devbox in running state, or the created devbox if waiting times out and + shutdown_on_timeout is False. Raises: PollingTimeout: If polling times out before devbox is running @@ -521,10 +524,16 @@ def create_and_await_running( idempotency_key=idempotency_key, ) - return self.await_running( - devbox.id, - polling_config=polling_config, - ) + try: + return self.await_running( + devbox.id, + polling_config=polling_config, + ) + except PollingTimeout: + if not shutdown_on_timeout: + return devbox + self.shutdown(devbox.id) + raise def list( self, @@ -2041,6 +2050,7 @@ async def create_and_await_running( mounts: Optional[Iterable[Mount]] | Omit = omit, name: Optional[str] | Omit = omit, polling_config: PollingConfig | None = None, + shutdown_on_timeout: bool = True, secrets: Optional[Dict[str, str]] | Omit = omit, snapshot_id: Optional[str] | Omit = omit, tunnel: Optional[devbox_create_params.Tunnel] | Omit = omit, @@ -2059,9 +2069,11 @@ async def create_and_await_running( Args: See the `create` method for detailed documentation. polling_config: Optional polling configuration + shutdown_on_timeout: Shutdown the created devbox if waiting for running state times out. Returns: - The devbox in running state + The devbox in running state, or the created devbox if waiting times out and + shutdown_on_timeout is False. Raises: PollingTimeout: If polling times out before devbox is running @@ -2092,10 +2104,16 @@ async def create_and_await_running( idempotency_key=idempotency_key, ) - return await self.await_running( - devbox.id, - polling_config=polling_config, - ) + try: + return await self.await_running( + devbox.id, + polling_config=polling_config, + ) + except PollingTimeout: + if not shutdown_on_timeout: + return devbox + await self.shutdown(devbox.id) + raise async def await_running( self, diff --git a/src/runloop_api_client/sdk/_types.py b/src/runloop_api_client/sdk/_types.py index 92070988e..a7a0c0e9b 100644 --- a/src/runloop_api_client/sdk/_types.py +++ b/src/runloop_api_client/sdk/_types.py @@ -92,11 +92,16 @@ class LongPollingRequestOptions(LongRequestOptions, PollingRequestOptions): # t pass -class SDKDevboxCreateParams(DevboxCreateParams, LongPollingRequestOptions): +class DevboxTimeoutCleanupOptions(TypedDict, total=False): + shutdown_on_timeout: bool + """Shutdown the created devbox when waiting for it to run times out. Defaults to true.""" + + +class SDKDevboxCreateParams(DevboxCreateParams, LongPollingRequestOptions, DevboxTimeoutCleanupOptions): pass -class SDKDevboxCreateFromImageParams(DevboxBaseCreateParams, LongPollingRequestOptions): +class SDKDevboxCreateFromImageParams(DevboxBaseCreateParams, LongPollingRequestOptions, DevboxTimeoutCleanupOptions): pass diff --git a/tests/api_resources/test_devboxes.py b/tests/api_resources/test_devboxes.py index 0d9495f80..714ad3435 100644 --- a/tests/api_resources/test_devboxes.py +++ b/tests/api_resources/test_devboxes.py @@ -4,7 +4,7 @@ import os from typing import Any, cast -from unittest.mock import Mock, patch +from unittest.mock import Mock, AsyncMock, patch import httpx import pytest @@ -1546,6 +1546,49 @@ def test_method_create_and_await_running_await_failure(self, client: Runloop) -> name="test", ) + @parametrize + def test_method_create_and_await_running_timeout_shuts_down(self, client: Runloop) -> None: + devbox = DevboxView( + id="test_id", + status="provisioning", + capabilities=[], + create_time_ms=1234567890, + launch_parameters=LaunchParameters(resource_size_request="X_SMALL"), + metadata={}, + state_transitions=[], + ) + timeout = PollingTimeout("Timed out", devbox) + + with patch.object(client.devboxes, "create", return_value=devbox): + with patch.object(client.devboxes, "await_running", side_effect=timeout): + with patch.object(client.devboxes, "shutdown") as mock_shutdown: + with pytest.raises(PollingTimeout) as exc_info: + client.devboxes.create_and_await_running() + + assert exc_info.value is timeout + mock_shutdown.assert_called_once_with("test_id") + + @parametrize + def test_method_create_and_await_running_timeout_returns_devbox_when_configured(self, client: Runloop) -> None: + devbox = DevboxView( + id="test_id", + status="provisioning", + capabilities=[], + create_time_ms=1234567890, + launch_parameters=LaunchParameters(resource_size_request="X_SMALL"), + metadata={}, + state_transitions=[], + ) + timeout = PollingTimeout("Timed out", devbox) + + with patch.object(client.devboxes, "create", return_value=devbox): + with patch.object(client.devboxes, "await_running", side_effect=timeout): + with patch.object(client.devboxes, "shutdown") as mock_shutdown: + result = client.devboxes.create_and_await_running(shutdown_on_timeout=False) + + assert result is devbox + mock_shutdown.assert_not_called() + @parametrize def test_method_await_suspended_success(self, client: Runloop) -> None: """Test await_suspended with successful polling to suspended state""" @@ -1744,6 +1787,51 @@ async def test_method_create(self, async_client: AsyncRunloop) -> None: devbox = await async_client.devboxes.create() assert_matches_type(DevboxView, devbox, path=["response"]) + @parametrize + async def test_method_create_and_await_running_timeout_shuts_down(self, async_client: AsyncRunloop) -> None: + devbox = DevboxView( + id="test_id", + status="provisioning", + capabilities=[], + create_time_ms=1234567890, + launch_parameters=LaunchParameters(resource_size_request="X_SMALL"), + metadata={}, + state_transitions=[], + ) + timeout = PollingTimeout("Timed out", devbox) + + with patch.object(async_client.devboxes, "create", AsyncMock(return_value=devbox)): + with patch.object(async_client.devboxes, "await_running", AsyncMock(side_effect=timeout)): + with patch.object(async_client.devboxes, "shutdown", AsyncMock()) as mock_shutdown: + with pytest.raises(PollingTimeout) as exc_info: + await async_client.devboxes.create_and_await_running() + + assert exc_info.value is timeout + mock_shutdown.assert_awaited_once_with("test_id") + + @parametrize + async def test_method_create_and_await_running_timeout_returns_devbox_when_configured( + self, async_client: AsyncRunloop + ) -> None: + devbox = DevboxView( + id="test_id", + status="provisioning", + capabilities=[], + create_time_ms=1234567890, + launch_parameters=LaunchParameters(resource_size_request="X_SMALL"), + metadata={}, + state_transitions=[], + ) + timeout = PollingTimeout("Timed out", devbox) + + with patch.object(async_client.devboxes, "create", AsyncMock(return_value=devbox)): + with patch.object(async_client.devboxes, "await_running", AsyncMock(side_effect=timeout)): + with patch.object(async_client.devboxes, "shutdown", AsyncMock()) as mock_shutdown: + result = await async_client.devboxes.create_and_await_running(shutdown_on_timeout=False) + + assert result is devbox + mock_shutdown.assert_not_awaited() + @parametrize async def test_method_create_with_all_params(self, async_client: AsyncRunloop) -> None: devbox = await async_client.devboxes.create( diff --git a/tests/sdk/test_async_ops.py b/tests/sdk/test_async_ops.py index 82a9535ce..ea26d5c44 100644 --- a/tests/sdk/test_async_ops.py +++ b/tests/sdk/test_async_ops.py @@ -65,11 +65,13 @@ async def test_create(self, mock_async_client: AsyncMock, devbox_view: MockDevbo name="test-devbox", metadata={"key": "value"}, polling_config=PollingConfig(timeout_seconds=60.0), + shutdown_on_timeout=False, ) assert isinstance(devbox, AsyncDevbox) assert devbox.id == "dbx_123" mock_async_client.devboxes.create_and_await_running.assert_awaited_once() + assert mock_async_client.devboxes.create_and_await_running.call_args.kwargs["shutdown_on_timeout"] is False @pytest.mark.asyncio async def test_create_from_blueprint_id(self, mock_async_client: AsyncMock, devbox_view: MockDevboxView) -> None: diff --git a/tests/sdk/test_ops.py b/tests/sdk/test_ops.py index 974c47e3d..64035c3bc 100644 --- a/tests/sdk/test_ops.py +++ b/tests/sdk/test_ops.py @@ -64,11 +64,13 @@ def test_create(self, mock_client: Mock, devbox_view: MockDevboxView) -> None: name="test-devbox", metadata={"key": "value"}, polling_config=PollingConfig(timeout_seconds=60.0), + shutdown_on_timeout=False, ) assert isinstance(devbox, Devbox) assert devbox.id == "dbx_123" mock_client.devboxes.create_and_await_running.assert_called_once() + assert mock_client.devboxes.create_and_await_running.call_args.kwargs["shutdown_on_timeout"] is False def test_create_from_blueprint_id(self, mock_client: Mock, devbox_view: MockDevboxView) -> None: """Test create_from_blueprint_id method."""