diff --git a/src/blaxel/core/sandbox/default/sandbox.py b/src/blaxel/core/sandbox/default/sandbox.py index 79831a6a..54e792e6 100644 --- a/src/blaxel/core/sandbox/default/sandbox.py +++ b/src/blaxel/core/sandbox/default/sandbox.py @@ -8,8 +8,12 @@ import httpx from ...client.api.compute.create_sandbox import asyncio as create_sandbox +from ...client.api.compute.create_sandbox_snapshot import asyncio as create_sandbox_snapshot from ...client.api.compute.delete_sandbox import asyncio as delete_sandbox +from ...client.api.compute.delete_sandbox_snapshot import asyncio as delete_sandbox_snapshot +from ...client.api.compute.fork_sandbox import asyncio as fork_sandbox from ...client.api.compute.get_sandbox import asyncio as get_sandbox +from ...client.api.compute.list_sandbox_snapshots import asyncio as list_sandbox_snapshots from ...client.api.compute.list_sandboxes import asyncio as list_sandboxes from ...client.api.compute.update_sandbox import asyncio as update_sandbox from ...client.client import client @@ -17,9 +21,13 @@ Metadata, MetadataLabels, Sandbox, + SandboxForkRequest, + SandboxForkResponse, SandboxLifecycle, SandboxRuntime, SandboxRuntimeExtraArgs, + SandboxSnapshot, + SandboxSnapshotRequest, SandboxSpec, ) from ...client.models import ( @@ -88,6 +96,23 @@ def _is_sandbox_not_found(error: SandboxAPIError) -> bool: return error.status_code == 404 or error.code in {404, "404"} +def _unwrap_response(response, action: str, *, allow_none: bool = False): + """Raise a SandboxAPIError for error/empty responses, else return the payload. + + Shared by the fork/snapshot helpers, which call generated control-plane API + functions that return ``Union[Error, T] | None``. Void endpoints (e.g. delete, + which returns 204 No Content) legitimately return ``None`` — pass + ``allow_none=True`` for those. + """ + if isinstance(response, Error): + status_code = response.code if response.code is not UNSET else None + message = response.message if response.message is not UNSET else response.error + raise SandboxAPIError(message, status_code=status_code, code=response.error) + if response is None and not allow_none: + raise SandboxAPIError(f"Failed to {action}") + return response + + def _sandbox_name( sandbox: Union[Sandbox, SandboxCreateConfiguration, Dict[str, Any]], ) -> str | None: @@ -209,6 +234,83 @@ async def fetch( """ return await self.network.fetch(port, path, method, **kwargs) + async def snapshot(self, name: str | None = None) -> SandboxSnapshot: + """Create a point-in-time snapshot of this sandbox. + + Snapshots capture the sandbox state and can be forked into new sandboxes + or applications. + + Args: + name: Optional human-readable name for the snapshot. + """ + body = SandboxSnapshotRequest(name=name) if name is not None else SandboxSnapshotRequest() + response = await create_sandbox_snapshot( + self.metadata.name, + client=client, + body=body, + ) + return _unwrap_response(response, "create snapshot") + + async def list_snapshots(self) -> list[SandboxSnapshot]: + """List the snapshots of this sandbox.""" + response = await list_sandbox_snapshots( + self.metadata.name, + client=client, + ) + return _unwrap_response(response, "list snapshots") + + async def delete_snapshot(self, snapshot_id: str) -> None: + """Delete a snapshot of this sandbox by its ID.""" + response = await delete_sandbox_snapshot( + self.metadata.name, + snapshot_id, + client=client, + ) + _unwrap_response(response, "delete snapshot", allow_none=True) + + async def fork( + self, + target_name: str, + *, + target_type: str = "sandbox", + port: int | None = None, + traffic: int | None = None, + custom_domain: str | None = None, + prefix: str | None = None, + snapshot_id: str | None = None, + ) -> SandboxForkResponse: + """Fork this sandbox into a new sandbox or application. + + Pass ``snapshot_id`` to fork from a specific snapshot (create a sandbox + from a snapshot) instead of the sandbox's live state. + + Args: + target_name: Name of the sandbox/application to create. + target_type: Resource type to fork into ("sandbox" or "application"). + port: Port to expose from the fork. + traffic: Canary traffic percentage (0-100) for an application fork. + custom_domain: Custom domain for an application fork. + prefix: URL prefix for an application fork. + snapshot_id: Snapshot ID to fork from. + """ + body = SandboxForkRequest(target_name=target_name, target_type=target_type) + if port is not None: + body.port = port + if traffic is not None: + body.traffic = traffic + if custom_domain is not None: + body.custom_domain = custom_domain + if prefix is not None: + body.prefix = prefix + if snapshot_id is not None: + body.snapshot_id = snapshot_id + response = await fork_sandbox( + self.metadata.name, + client=client, + body=body, + ) + return _unwrap_response(response, "fork sandbox") + async def wait(self, max_wait: int = 60000, interval: int = 1000) -> "SandboxInstance": logger.warning( "⚠️ Warning: sandbox.wait() is deprecated. You don't need to wait for the sandbox to be deployed anymore." diff --git a/src/blaxel/core/sandbox/sync/sandbox.py b/src/blaxel/core/sandbox/sync/sandbox.py index a6e2cc05..9750678e 100644 --- a/src/blaxel/core/sandbox/sync/sandbox.py +++ b/src/blaxel/core/sandbox/sync/sandbox.py @@ -7,17 +7,25 @@ import httpx from ...client.api.compute.create_sandbox import sync as create_sandbox +from ...client.api.compute.create_sandbox_snapshot import sync as create_sandbox_snapshot from ...client.api.compute.delete_sandbox import sync as delete_sandbox +from ...client.api.compute.delete_sandbox_snapshot import sync as delete_sandbox_snapshot +from ...client.api.compute.fork_sandbox import sync as fork_sandbox from ...client.api.compute.get_sandbox import sync as get_sandbox +from ...client.api.compute.list_sandbox_snapshots import sync as list_sandbox_snapshots from ...client.api.compute.list_sandboxes import sync as list_sandboxes from ...client.api.compute.update_sandbox import sync as update_sandbox from ...client.client import client from ...client.models import ( Metadata, Sandbox, + SandboxForkRequest, + SandboxForkResponse, SandboxLifecycle, SandboxRuntime, SandboxRuntimeExtraArgs, + SandboxSnapshot, + SandboxSnapshotRequest, SandboxSpec, ) from ...client.models import ( @@ -38,6 +46,7 @@ _is_sandbox_conflict, _is_sandbox_not_found, _sandbox_name, + _unwrap_response, ) from ..types import ( SandboxConfiguration, @@ -141,6 +150,83 @@ def fetch(self, port: int, path: str = "/", method: str = "GET", **kwargs) -> "h """ return self.network.fetch(port, path, method, **kwargs) + def snapshot(self, name: str | None = None) -> SandboxSnapshot: + """Create a point-in-time snapshot of this sandbox. + + Snapshots capture the sandbox state and can be forked into new sandboxes + or applications. + + Args: + name: Optional human-readable name for the snapshot. + """ + body = SandboxSnapshotRequest(name=name) if name is not None else SandboxSnapshotRequest() + response = create_sandbox_snapshot( + self.metadata.name, + client=client, + body=body, + ) + return _unwrap_response(response, "create snapshot") + + def list_snapshots(self) -> list[SandboxSnapshot]: + """List the snapshots of this sandbox.""" + response = list_sandbox_snapshots( + self.metadata.name, + client=client, + ) + return _unwrap_response(response, "list snapshots") + + def delete_snapshot(self, snapshot_id: str) -> None: + """Delete a snapshot of this sandbox by its ID.""" + response = delete_sandbox_snapshot( + self.metadata.name, + snapshot_id, + client=client, + ) + _unwrap_response(response, "delete snapshot", allow_none=True) + + def fork( + self, + target_name: str, + *, + target_type: str = "sandbox", + port: int | None = None, + traffic: int | None = None, + custom_domain: str | None = None, + prefix: str | None = None, + snapshot_id: str | None = None, + ) -> SandboxForkResponse: + """Fork this sandbox into a new sandbox or application. + + Pass ``snapshot_id`` to fork from a specific snapshot (create a sandbox + from a snapshot) instead of the sandbox's live state. + + Args: + target_name: Name of the sandbox/application to create. + target_type: Resource type to fork into ("sandbox" or "application"). + port: Port to expose from the fork. + traffic: Canary traffic percentage (0-100) for an application fork. + custom_domain: Custom domain for an application fork. + prefix: URL prefix for an application fork. + snapshot_id: Snapshot ID to fork from. + """ + body = SandboxForkRequest(target_name=target_name, target_type=target_type) + if port is not None: + body.port = port + if traffic is not None: + body.traffic = traffic + if custom_domain is not None: + body.custom_domain = custom_domain + if prefix is not None: + body.prefix = prefix + if snapshot_id is not None: + body.snapshot_id = snapshot_id + response = fork_sandbox( + self.metadata.name, + client=client, + body=body, + ) + return _unwrap_response(response, "fork sandbox") + def wait(self, max_wait: int = 60000, interval: int = 1000) -> "SyncSandboxInstance": logger.warning( "⚠️ Warning: sandbox.wait() is deprecated. You don't need to wait for the sandbox to be deployed anymore." diff --git a/tests/core/test_sandbox.py b/tests/core/test_sandbox.py index 4c410615..5540f4ce 100644 --- a/tests/core/test_sandbox.py +++ b/tests/core/test_sandbox.py @@ -834,3 +834,112 @@ def test_sync_create_sends_name_when_provided(): body = mock_create_sandbox.call_args.kwargs["body"] assert _body_metadata_name(body) == "mysbx" + + +@pytest.mark.asyncio +async def test_fork_defaults_to_sandbox_target(): + sandbox = sandbox_instance("my-sandbox") + + with patch( + "blaxel.core.sandbox.default.sandbox.fork_sandbox", new_callable=AsyncMock + ) as mock_fork: + mock_fork.return_value = MagicMock(name="my-sandbox-copy", type="sandbox") + + await sandbox.fork("my-sandbox-copy") + + assert mock_fork.call_args.args[0] == "my-sandbox" + body = mock_fork.call_args.kwargs["body"] + assert body.target_name == "my-sandbox-copy" + assert body.target_type == "sandbox" + + +@pytest.mark.asyncio +async def test_fork_forwards_application_options_and_snapshot(): + sandbox = sandbox_instance("my-sandbox") + + with patch( + "blaxel.core.sandbox.default.sandbox.fork_sandbox", new_callable=AsyncMock + ) as mock_fork: + mock_fork.return_value = MagicMock() + + await sandbox.fork( + "my-app", + target_type="application", + traffic=100, + port=8080, + custom_domain="app.example.com", + snapshot_id="snap_abc123", + ) + + body = mock_fork.call_args.kwargs["body"] + assert body.target_name == "my-app" + assert body.target_type == "application" + assert body.traffic == 100 + assert body.port == 8080 + assert body.custom_domain == "app.example.com" + assert body.snapshot_id == "snap_abc123" + + +@pytest.mark.asyncio +async def test_snapshot_sends_optional_name(): + sandbox = sandbox_instance("my-sandbox") + + with patch( + "blaxel.core.sandbox.default.sandbox.create_sandbox_snapshot", new_callable=AsyncMock + ) as mock_snapshot: + mock_snapshot.return_value = MagicMock() + + await sandbox.snapshot("before") + + assert mock_snapshot.call_args.args[0] == "my-sandbox" + assert mock_snapshot.call_args.kwargs["body"].name == "before" + + +@pytest.mark.asyncio +async def test_delete_snapshot_accepts_none_204_response(): + sandbox = sandbox_instance("my-sandbox") + + with patch( + "blaxel.core.sandbox.default.sandbox.delete_sandbox_snapshot", new_callable=AsyncMock + ) as mock_delete: + # Generated client returns None for a successful 204 No Content. + mock_delete.return_value = None + + await sandbox.delete_snapshot("snap_abc123") + + assert mock_delete.call_args.args == ("my-sandbox", "snap_abc123") + + +@pytest.mark.asyncio +async def test_fork_raises_on_error_response(): + from blaxel.core.client.models.error import Error + + sandbox = sandbox_instance("my-sandbox") + + with patch( + "blaxel.core.sandbox.default.sandbox.fork_sandbox", new_callable=AsyncMock + ) as mock_fork: + mock_fork.return_value = Error(error="boom", code=400) + + with pytest.raises(SandboxAPIError): + await sandbox.fork("my-sandbox-copy") + + +def test_sync_fork_and_snapshot_helpers(): + sandbox = sandbox_instance("my-sandbox", cls=SyncSandboxInstance) + + with ( + patch("blaxel.core.sandbox.sync.sandbox.fork_sandbox") as mock_fork, + patch("blaxel.core.sandbox.sync.sandbox.create_sandbox_snapshot") as mock_snapshot, + ): + mock_fork.return_value = MagicMock() + mock_snapshot.return_value = MagicMock() + + sandbox.fork("my-sandbox-copy", snapshot_id="snap_abc123") + sandbox.snapshot("before") + + fork_body = mock_fork.call_args.kwargs["body"] + assert fork_body.target_name == "my-sandbox-copy" + assert fork_body.target_type == "sandbox" + assert fork_body.snapshot_id == "snap_abc123" + assert mock_snapshot.call_args.kwargs["body"].name == "before"