diff --git a/src/blaxel/core/sandbox/__init__.py b/src/blaxel/core/sandbox/__init__.py index 02c217f..38e952d 100644 --- a/src/blaxel/core/sandbox/__init__.py +++ b/src/blaxel/core/sandbox/__init__.py @@ -38,6 +38,7 @@ SandboxConfiguration, SandboxCreateConfiguration, SandboxFilesystemFile, + SandboxUpdateNetwork, SessionCreateOptions, SessionWithToken, StreamHandle, @@ -59,6 +60,7 @@ "AsyncStreamHandle", "SandboxFilesystemFile", "CopyResponse", + "SandboxUpdateNetwork", "Sandbox", "SandboxFileSystem", "SandboxPreviews", diff --git a/src/blaxel/core/sandbox/default/sandbox.py b/src/blaxel/core/sandbox/default/sandbox.py index a069c67..f80a600 100644 --- a/src/blaxel/core/sandbox/default/sandbox.py +++ b/src/blaxel/core/sandbox/default/sandbox.py @@ -32,6 +32,7 @@ SandboxConfiguration, SandboxCreateConfiguration, SandboxUpdateMetadata, + SandboxUpdateNetwork, SessionWithToken, ) from .codegen import SandboxCodegen @@ -477,6 +478,42 @@ async def update_lifecycle( return cls(response) + @classmethod + async def update_network( + cls, sandbox_name: str, network: SandboxUpdateNetwork + ) -> "SandboxInstance": + """Update sandbox network configuration without recreating it. + + Args: + sandbox_name: The name of the sandbox to update + network: The new network configuration + + Returns: + A new SandboxInstance with updated network configuration + """ + sandbox_instance = await cls.get(sandbox_name) + sandbox = sandbox_instance.sandbox + + updated_sandbox = Sandbox.from_dict(sandbox.to_dict()) + if updated_sandbox.spec is None: + raise ValueError(f"Sandbox {sandbox_name} has invalid spec") + + if network.network is not None: + if isinstance(network.network, dict): + updated_sandbox.spec.network = SandboxNetworkModel.from_dict(network.network) + else: + updated_sandbox.spec.network = network.network + else: + updated_sandbox.spec.network = UNSET + + response = await update_sandbox( + sandbox_name=sandbox_name, + client=client, + body=updated_sandbox, + ) + + return cls(response) + @classmethod async def create_if_not_exists( cls, sandbox: Union[Sandbox, SandboxCreateConfiguration, Dict[str, Any]] diff --git a/src/blaxel/core/sandbox/sync/sandbox.py b/src/blaxel/core/sandbox/sync/sandbox.py index 9a76f50..cf64e2a 100644 --- a/src/blaxel/core/sandbox/sync/sandbox.py +++ b/src/blaxel/core/sandbox/sync/sandbox.py @@ -37,6 +37,7 @@ SandboxConfiguration, SandboxCreateConfiguration, SandboxUpdateMetadata, + SandboxUpdateNetwork, SessionWithToken, ) from .codegen import SyncSandboxCodegen @@ -392,6 +393,42 @@ def update_lifecycle( return cls(response) + @classmethod + def update_network( + cls, sandbox_name: str, network: SandboxUpdateNetwork + ) -> "SyncSandboxInstance": + """Update sandbox network configuration without recreating it. + + Args: + sandbox_name: The name of the sandbox to update + network: The new network configuration + + Returns: + A new SyncSandboxInstance with updated network configuration + """ + sandbox_instance = cls.get(sandbox_name) + sandbox = sandbox_instance.sandbox + + updated_sandbox = Sandbox.from_dict(sandbox.to_dict()) + if updated_sandbox.spec is None: + raise ValueError(f"Sandbox {sandbox_name} has invalid spec") + + if network.network is not None: + if isinstance(network.network, dict): + updated_sandbox.spec.network = SandboxNetworkModel.from_dict(network.network) + else: + updated_sandbox.spec.network = network.network + else: + updated_sandbox.spec.network = UNSET + + response = update_sandbox( + sandbox_name=sandbox_name, + client=client, + body=updated_sandbox, + ) + + return cls(response) + @classmethod def create_if_not_exists( cls, sandbox: Union[Sandbox, SandboxCreateConfiguration, Dict[str, Any]] diff --git a/src/blaxel/core/sandbox/types.py b/src/blaxel/core/sandbox/types.py index b43ed1f..22628d6 100644 --- a/src/blaxel/core/sandbox/types.py +++ b/src/blaxel/core/sandbox/types.py @@ -146,6 +146,16 @@ def __init__( self.display_name = display_name +class SandboxUpdateNetwork: + """Configuration for updating sandbox network configuration.""" + + def __init__( + self, + network: Union[SandboxNetwork, Dict[str, Any]] | None = None, + ): + self.network = network + + class SandboxCreateConfiguration: """Simplified configuration for creating sandboxes with default values.""" diff --git a/tests/integration/core/sandbox/test_network.py b/tests/integration/core/sandbox/test_network.py new file mode 100644 index 0000000..6a06f4e --- /dev/null +++ b/tests/integration/core/sandbox/test_network.py @@ -0,0 +1,111 @@ +import pytest + +from blaxel.core import SandboxInstance +from blaxel.core.sandbox.types import SandboxUpdateNetwork +from tests.helpers import ( + default_image, + default_labels, + unique_name, +) + +# ============================================================================= +# Sandbox UpdateNetwork Tests +# ============================================================================= + + +@pytest.mark.asyncio(loop_scope="class") +class TestSandboxUpdateNetwork: + """Test sandbox update_network operations.""" + + async def test_updates_sandbox_network_with_allowed_domains(self): + """Test updating sandbox network with allowed domains using httpbin.""" + name = unique_name("update-net-allow") + await SandboxInstance.create( + { + "name": name, + "image": default_image, + "labels": default_labels, + } + ) + + try: + updated = await SandboxInstance.update_network( + name, + SandboxUpdateNetwork(network={"allowedDomains": ["httpbin.org", "*.httpbin.org"]}), + ) + assert updated.spec.network is not None + assert set(updated.spec.network.allowed_domains) == {"httpbin.org", "*.httpbin.org"} + finally: + await SandboxInstance.delete(name) + + async def test_updates_sandbox_network_with_forbidden_domains(self): + """Test updating sandbox network with forbidden domains using httpbin.""" + name = unique_name("update-net-forbid") + await SandboxInstance.create( + { + "name": name, + "image": default_image, + "labels": default_labels, + } + ) + + try: + updated = await SandboxInstance.update_network( + name, + SandboxUpdateNetwork( + network={"forbiddenDomains": ["httpbin.org", "*.httpbin.org"]} + ), + ) + assert updated.spec.network is not None + assert set(updated.spec.network.forbidden_domains) == {"httpbin.org", "*.httpbin.org"} + finally: + await SandboxInstance.delete(name) + + async def test_updates_sandbox_network_with_model_object(self): + """Test updating sandbox network using SandboxNetwork model object.""" + from blaxel.core.client.models import SandboxNetwork + + name = unique_name("update-net-model") + await SandboxInstance.create( + { + "name": name, + "image": default_image, + "labels": default_labels, + } + ) + + try: + network_config = SandboxNetwork( + allowed_domains=["httpbin.org"], + ) + updated = await SandboxInstance.update_network( + name, + SandboxUpdateNetwork(network=network_config), + ) + assert updated.spec.network is not None + assert updated.spec.network.allowed_domains == ["httpbin.org"] + finally: + await SandboxInstance.delete(name) + + async def test_clears_sandbox_network_config(self): + """Test clearing sandbox network configuration by passing network=None.""" + name = unique_name("update-net-clear") + await SandboxInstance.create( + { + "name": name, + "image": default_image, + "network": {"allowedDomains": ["httpbin.org"]}, + "labels": default_labels, + } + ) + + try: + updated = await SandboxInstance.update_network( + name, + SandboxUpdateNetwork(network=None), + ) + from blaxel.core.client.types import UNSET + + assert updated.spec.network is UNSET or updated.spec.network is None + finally: + await SandboxInstance.delete(name)