wehub-resource-sync
136 lines
3.9 KiB
Python
136 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from agents.sandbox.errors import WorkspaceArchiveReadError
|
|
from agents.sandbox.session.mount_lifecycle import with_ephemeral_mounts_removed
|
|
|
|
|
|
class _FakeMountStrategy:
|
|
def __init__(
|
|
self,
|
|
events: list[str],
|
|
*,
|
|
name: str,
|
|
fail_teardown: bool = False,
|
|
fail_restore: bool = False,
|
|
) -> None:
|
|
self._events = events
|
|
self._name = name
|
|
self._fail_teardown = fail_teardown
|
|
self._fail_restore = fail_restore
|
|
|
|
async def teardown_for_snapshot(
|
|
self,
|
|
mount: object,
|
|
session: object,
|
|
path: Path,
|
|
) -> None:
|
|
_ = (mount, session, path)
|
|
self._events.append(f"teardown:{self._name}")
|
|
if self._fail_teardown:
|
|
raise RuntimeError(f"teardown failed: {self._name}")
|
|
|
|
async def restore_after_snapshot(
|
|
self,
|
|
mount: object,
|
|
session: object,
|
|
path: Path,
|
|
) -> None:
|
|
_ = (mount, session, path)
|
|
self._events.append(f"restore:{self._name}")
|
|
if self._fail_restore:
|
|
raise RuntimeError(f"restore failed: {self._name}")
|
|
|
|
|
|
class _FakeMount:
|
|
def __init__(self, strategy: _FakeMountStrategy) -> None:
|
|
self.mount_strategy = strategy
|
|
|
|
|
|
class _FakeManifest:
|
|
def __init__(self, mounts: list[tuple[_FakeMount, Path]]) -> None:
|
|
self._mounts = mounts
|
|
|
|
def ephemeral_mount_targets(self) -> list[tuple[_FakeMount, Path]]:
|
|
return self._mounts
|
|
|
|
|
|
class _FakeState:
|
|
def __init__(self, manifest: _FakeManifest) -> None:
|
|
self.manifest = manifest
|
|
|
|
|
|
class _FakeSession:
|
|
def __init__(self, manifest: _FakeManifest) -> None:
|
|
self.state = _FakeState(manifest)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_with_ephemeral_mounts_removed_restores_in_reverse_order() -> None:
|
|
events: list[str] = []
|
|
left = _FakeMount(_FakeMountStrategy(events, name="left"))
|
|
right = _FakeMount(_FakeMountStrategy(events, name="right"))
|
|
session = _FakeSession(
|
|
_FakeManifest(
|
|
[
|
|
(left, Path("/workspace/left")),
|
|
(right, Path("/workspace/right")),
|
|
]
|
|
)
|
|
)
|
|
|
|
async def operation() -> str:
|
|
events.append("operation")
|
|
return "persisted"
|
|
|
|
result = await with_ephemeral_mounts_removed(
|
|
cast(Any, session),
|
|
operation,
|
|
error_path=Path("/workspace"),
|
|
error_cls=WorkspaceArchiveReadError,
|
|
operation_error_context_key="snapshot_error_before_remount_corruption",
|
|
)
|
|
|
|
assert result == "persisted"
|
|
assert events == [
|
|
"teardown:left",
|
|
"teardown:right",
|
|
"operation",
|
|
"restore:right",
|
|
"restore:left",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_with_ephemeral_mounts_removed_reports_restore_error_after_operation_error() -> None:
|
|
events: list[str] = []
|
|
mount = _FakeMount(_FakeMountStrategy(events, name="mount", fail_restore=True))
|
|
session = _FakeSession(_FakeManifest([(mount, Path("/workspace/mount"))]))
|
|
operation_error = WorkspaceArchiveReadError(
|
|
path=Path("/workspace"),
|
|
context={"reason": "persist_failed"},
|
|
)
|
|
|
|
async def operation() -> bytes:
|
|
events.append("operation")
|
|
raise operation_error
|
|
|
|
with pytest.raises(WorkspaceArchiveReadError) as exc_info:
|
|
await with_ephemeral_mounts_removed(
|
|
cast(Any, session),
|
|
operation,
|
|
error_path=Path("/workspace"),
|
|
error_cls=WorkspaceArchiveReadError,
|
|
operation_error_context_key="snapshot_error_before_remount_corruption",
|
|
)
|
|
|
|
assert events == ["teardown:mount", "operation", "restore:mount"]
|
|
assert exc_info.value.context["snapshot_error_before_remount_corruption"] == {
|
|
"message": operation_error.message,
|
|
}
|
|
assert isinstance(exc_info.value.cause, RuntimeError)
|