tests/test_vfs_atomic_rollback.py

from __future__ import annotations

import asyncio
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from agents.unified_loop import UnifiedAgentLoop
from tools.registry import _delete_file, _read_file, _write_file


class FakeExecutor:
    def __init__(self, failures: set[str] | None = None) -> None:
        self.calls: list[tuple[str, dict]] = []
        self.failures = failures or set()

    async def run_tool(self, name: str, inputs: dict, timeout: float = 30.0) -> dict:
        self.calls.append((name, inputs))
        if name in self.failures:
            return {"success": False, "error": f"forced failure: {name}", "output": None}
        return {"success": True, "output": {"ok": True}}


class VfsAtomicRollbackTests(unittest.IsolatedAsyncioTestCase):
    def make_loop(self, executor: FakeExecutor) -> UnifiedAgentLoop:
        loop = UnifiedAgentLoop.__new__(UnifiedAgentLoop)
        loop.executor = executor
        loop._write_snapshots = {}
        return loop

    async def test_rollback_restores_existing_and_deletes_new_files(self) -> None:
        executor = FakeExecutor()
        loop = self.make_loop(executor)
        loop._write_snapshots = {"existing.txt": "before", "created.txt": None}

        await loop._rollback_writes()

        self.assertEqual(
            [(name, inputs) for name, inputs in executor.calls],
            [
                ("write_file", {"path": "existing.txt", "content": "before"}),
                ("delete_file", {"path": "created.txt"}),
            ],
        )
        self.assertEqual(loop._write_snapshots, {})

    async def test_failed_restore_is_not_marked_clean(self) -> None:
        executor = FakeExecutor({"write_file"})
        loop = self.make_loop(executor)
        loop._write_snapshots = {"existing.txt": "before", "created.txt": None}

        with self.assertRaisesRegex(RuntimeError, "rollback incompleto"):
            await loop._rollback_writes()

        self.assertEqual(loop._write_snapshots, {"existing.txt": "before"})
        self.assertEqual(executor.calls[1][0], "delete_file")

    async def test_delete_file_respects_fs_jail_and_is_idempotent(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            with patch.dict("os.environ", {"FS_TOOL_ROOT": root}, clear=False):
                target = Path(root) / "created.txt"
                target.write_text("created", encoding="utf-8")
                deleted = await _delete_file("created.txt")
                repeated = await _delete_file("created.txt")
                outside = await _delete_file("../outside.txt")

        self.assertTrue(deleted["ok"])
        self.assertTrue(deleted["deleted"])
        self.assertTrue(repeated["ok"])
        self.assertFalse(repeated["deleted"])
        self.assertFalse(outside["ok"])


if __name__ == "__main__":
    unittest.main()