91e75e620b
CI: cua-driver distro-compat matrix / debian:12 (glibc 2.36) (push) Has been cancelled
CI: SPDX Headers / Check SPDX headers (warn-only) (push) Has been cancelled
CD: Docs MCP Server / build (linux/amd64) (push) Has been cancelled
CD: Docs MCP Server / build (linux/arm64) (push) Has been cancelled
CD: Docs MCP Server / merge (push) Has been cancelled
CI: cua-driver distro-compat matrix / Resolve release version (push) Has been cancelled
CI: cua-driver distro-compat matrix / fedora:41 (glibc 2.40) (push) Has been cancelled
CI: cua-driver distro-compat matrix / rockylinux:9 (glibc 2.34) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:22.04 (glibc 2.35) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:24.04 (glibc 2.39) (push) Has been cancelled
CI: cua-driver distro-compat matrix / Distro compat summary (push) Has been cancelled
CI: Rust Linux unit / Rust Linux unit and compile (push) Has been cancelled
CI: Rust Windows unit / Rust Windows unit and compile (push) Has been cancelled
CI: Nix Linux Rust source / Nix / compositor build (push) Has been cancelled
CI: Nix Linux Rust source / Nix / driver package (push) Has been cancelled
CI: Nix Linux Rust source / Nix / Rust unit tests (push) Has been cancelled
337 lines
12 KiB
Python
337 lines
12 KiB
Python
"""Unit tests for VM cleanup on connection failure and destroy() resilience.
|
|
|
|
These tests mock CloudTransport so they run without a real cloud API.
|
|
They verify that:
|
|
1. _create() cleans up a provisioned VM when _connect() fails.
|
|
2. destroy() runs every cleanup step independently — a failure in one
|
|
does not prevent the others from executing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from cua_sandbox.image import Image
|
|
from cua_sandbox.sandbox import Sandbox
|
|
from cua_sandbox.transport.cloud import CloudTransport
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_cloud_transport(*, name: str = "test-vm") -> CloudTransport:
|
|
"""Return a CloudTransport with internal state set as if _create_vm() succeeded."""
|
|
t = CloudTransport.__new__(CloudTransport)
|
|
t._name = name
|
|
t._api_key_override = "sk-fake"
|
|
t._base_url = "https://api.example.com"
|
|
t._image = None
|
|
t._cpu = None
|
|
t._memory_mb = None
|
|
t._disk_gb = None
|
|
t._region = "us-east-1"
|
|
t._inner = None
|
|
t._api_client = None
|
|
return t
|
|
|
|
|
|
def _make_sandbox(transport: CloudTransport, **kwargs) -> Sandbox:
|
|
"""Return a Sandbox wrapping *transport* without calling _connect()."""
|
|
return Sandbox(
|
|
transport,
|
|
name=transport._name,
|
|
_ephemeral=kwargs.get("ephemeral", True),
|
|
_telemetry_enabled=False,
|
|
)
|
|
|
|
|
|
# ===================================================================
|
|
# 1. _create() cleans up on _connect() failure
|
|
# ===================================================================
|
|
|
|
|
|
class TestCreateCleansUpOnConnectFailure:
|
|
"""Sandbox._create() should delete the cloud VM when _connect() raises."""
|
|
|
|
async def test_delete_vm_called_on_timeout(self):
|
|
"""ReadTimeout during _connect() triggers delete_vm()."""
|
|
transport = _make_cloud_transport(name="orphan-vm")
|
|
transport.connect = AsyncMock(side_effect=httpx.ReadTimeout("poll timed out"))
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with patch(
|
|
"cua_sandbox.sandbox.CloudTransport",
|
|
return_value=transport,
|
|
):
|
|
with pytest.raises(httpx.ReadTimeout):
|
|
await Sandbox._create(
|
|
image=Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_delete_vm_called_on_generic_exception(self):
|
|
"""Any exception during _connect() triggers delete_vm()."""
|
|
transport = _make_cloud_transport(name="orphan-vm")
|
|
transport.connect = AsyncMock(side_effect=RuntimeError("unexpected"))
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with patch(
|
|
"cua_sandbox.sandbox.CloudTransport",
|
|
return_value=transport,
|
|
):
|
|
with pytest.raises(RuntimeError, match="unexpected"):
|
|
await Sandbox._create(
|
|
image=Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_original_exception_propagates_even_if_delete_fails(self):
|
|
"""The original connect error is re-raised even when delete_vm also fails."""
|
|
transport = _make_cloud_transport(name="orphan-vm")
|
|
transport.connect = AsyncMock(side_effect=TimeoutError("poll timeout"))
|
|
transport.delete_vm = AsyncMock(side_effect=httpx.ConnectError("api down"))
|
|
|
|
with patch(
|
|
"cua_sandbox.sandbox.CloudTransport",
|
|
return_value=transport,
|
|
):
|
|
with pytest.raises(TimeoutError, match="poll timeout"):
|
|
await Sandbox._create(
|
|
image=Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_no_cleanup_when_vm_not_yet_created(self):
|
|
"""If connect() fails before _create_vm (no _name), skip delete_vm."""
|
|
transport = _make_cloud_transport()
|
|
# Simulate: connect() fails before _create_vm sets _name
|
|
transport._name = None
|
|
transport.connect = AsyncMock(side_effect=ValueError("no api key"))
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with patch(
|
|
"cua_sandbox.sandbox.CloudTransport",
|
|
return_value=transport,
|
|
):
|
|
with pytest.raises(ValueError, match="no api key"):
|
|
await Sandbox._create(
|
|
image=Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
transport.delete_vm.assert_not_awaited()
|
|
|
|
async def test_keyboard_interrupt_still_cleans_up(self):
|
|
"""BaseException subclasses (KeyboardInterrupt) also trigger cleanup."""
|
|
transport = _make_cloud_transport(name="interrupted-vm")
|
|
transport.connect = AsyncMock(side_effect=KeyboardInterrupt)
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with patch(
|
|
"cua_sandbox.sandbox.CloudTransport",
|
|
return_value=transport,
|
|
):
|
|
with pytest.raises(KeyboardInterrupt):
|
|
await Sandbox._create(
|
|
image=Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
|
|
# ===================================================================
|
|
# 2. destroy() is resilient to individual step failures
|
|
# ===================================================================
|
|
|
|
|
|
class TestDestroyResilience:
|
|
"""Each cleanup step in destroy() should run independently."""
|
|
|
|
async def test_delete_vm_runs_even_if_disconnect_fails(self):
|
|
"""A failing disconnect() must not prevent delete_vm()."""
|
|
transport = _make_cloud_transport(name="leaky-vm")
|
|
transport.disconnect = AsyncMock(side_effect=OSError("connection reset"))
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
sb = _make_sandbox(transport)
|
|
await sb.destroy()
|
|
|
|
transport.disconnect.assert_awaited_once()
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_runtime_stop_runs_even_if_delete_vm_fails(self):
|
|
"""A failing delete_vm() must not prevent runtime cleanup."""
|
|
transport = _make_cloud_transport(name="leaky-vm")
|
|
transport.disconnect = AsyncMock()
|
|
transport.delete_vm = AsyncMock(side_effect=httpx.ConnectError("api down"))
|
|
|
|
runtime = AsyncMock()
|
|
runtime_info = MagicMock()
|
|
runtime_info.name = "leaky-vm"
|
|
|
|
sb = _make_sandbox(transport)
|
|
sb._runtime = runtime
|
|
sb._runtime_info = runtime_info
|
|
|
|
await sb.destroy()
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
runtime.delete.assert_awaited_once_with("leaky-vm")
|
|
|
|
async def test_destroy_succeeds_when_all_steps_fail(self):
|
|
"""destroy() must not raise even if every cleanup step fails."""
|
|
transport = _make_cloud_transport(name="total-fail")
|
|
transport.disconnect = AsyncMock(side_effect=OSError("disconnect fail"))
|
|
transport.delete_vm = AsyncMock(side_effect=httpx.ReadTimeout("delete fail"))
|
|
|
|
runtime = AsyncMock()
|
|
runtime.delete = AsyncMock(side_effect=RuntimeError("runtime fail"))
|
|
runtime_info = MagicMock()
|
|
runtime_info.name = "total-fail"
|
|
|
|
sb = _make_sandbox(transport)
|
|
sb._runtime = runtime
|
|
sb._runtime_info = runtime_info
|
|
|
|
# Should not raise
|
|
await sb.destroy()
|
|
|
|
transport.disconnect.assert_awaited_once()
|
|
transport.delete_vm.assert_awaited_once()
|
|
runtime.delete.assert_awaited_once()
|
|
|
|
async def test_destroy_happy_path(self):
|
|
"""All steps succeed — basic smoke test."""
|
|
transport = _make_cloud_transport(name="good-vm")
|
|
transport.disconnect = AsyncMock()
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
sb = _make_sandbox(transport)
|
|
await sb.destroy()
|
|
|
|
transport.disconnect.assert_awaited_once()
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_non_cloud_transport_skips_delete_vm(self):
|
|
"""Non-CloudTransport sandboxes should not call delete_vm."""
|
|
transport = AsyncMock() # generic mock, not a CloudTransport instance
|
|
sb = Sandbox(transport, name="local-vm", _ephemeral=True, _telemetry_enabled=False)
|
|
|
|
await sb.destroy()
|
|
|
|
transport.disconnect.assert_awaited_once()
|
|
# delete_vm should not be called since transport is not CloudTransport
|
|
assert not hasattr(transport, "delete_vm") or not transport.delete_vm.called
|
|
|
|
|
|
# ===================================================================
|
|
# 3. ephemeral() integration — cleanup through the context manager
|
|
# ===================================================================
|
|
|
|
|
|
class TestEphemeralCleanup:
|
|
"""Sandbox.ephemeral() should destroy the VM on normal and error exits."""
|
|
|
|
async def test_ephemeral_destroys_on_normal_exit(self):
|
|
"""VM is destroyed when the async-with block exits normally."""
|
|
transport = _make_cloud_transport(name="eph-ok")
|
|
transport.connect = AsyncMock()
|
|
transport.disconnect = AsyncMock()
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with (
|
|
patch.object(
|
|
CloudTransport,
|
|
"__init__",
|
|
lambda self, **kw: None,
|
|
),
|
|
patch.object(
|
|
CloudTransport,
|
|
"__new__",
|
|
lambda cls, **kw: transport,
|
|
),
|
|
):
|
|
async with Sandbox.ephemeral(
|
|
Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
) as sb:
|
|
assert sb.name == "eph-ok"
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_ephemeral_destroys_on_test_failure(self):
|
|
"""VM is destroyed even when the body raises an assertion error."""
|
|
transport = _make_cloud_transport(name="eph-fail")
|
|
transport.connect = AsyncMock()
|
|
transport.disconnect = AsyncMock()
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with (
|
|
patch.object(
|
|
CloudTransport,
|
|
"__init__",
|
|
lambda self, **kw: None,
|
|
),
|
|
patch.object(
|
|
CloudTransport,
|
|
"__new__",
|
|
lambda cls, **kw: transport,
|
|
),
|
|
):
|
|
with pytest.raises(AssertionError):
|
|
async with Sandbox.ephemeral(
|
|
Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
) as _sb:
|
|
raise AssertionError("test failed")
|
|
|
|
transport.delete_vm.assert_awaited_once()
|
|
|
|
async def test_ephemeral_cleans_up_when_create_connect_fails(self):
|
|
"""If _create raises after VM provisioning, the VM is still cleaned up."""
|
|
transport = _make_cloud_transport(name="eph-connect-fail")
|
|
transport.connect = AsyncMock(side_effect=httpx.ReadTimeout("poll timed out"))
|
|
transport.delete_vm = AsyncMock()
|
|
|
|
with (
|
|
patch.object(
|
|
CloudTransport,
|
|
"__init__",
|
|
lambda self, **kw: None,
|
|
),
|
|
patch.object(
|
|
CloudTransport,
|
|
"__new__",
|
|
lambda cls, **kw: transport,
|
|
),
|
|
):
|
|
with pytest.raises(httpx.ReadTimeout):
|
|
async with Sandbox.ephemeral(
|
|
Image.linux("ubuntu", "24.04"),
|
|
api_key="sk-fake",
|
|
telemetry_enabled=False,
|
|
) as _sb:
|
|
pass # never reached
|
|
|
|
transport.delete_vm.assert_awaited_once()
|