cleveragents-core/robot/helper_sandbox_integration.py

"""Helper utilities for sandbox infrastructure Robot integration tests.

Covers the integration between:
- SandboxStatus and transition validation (protocol.py)
- NoSandbox implementation (no_sandbox.py)
- SandboxFactory (factory.py)
- SandboxManager (manager.py)
- Merge strategies (merge.py)
"""

from __future__ import annotations

import json
import sys
import tempfile

from cleveragents.infrastructure.sandbox.factory import SandboxFactory
from cleveragents.infrastructure.sandbox.manager import SandboxManager
from cleveragents.infrastructure.sandbox.merge import (
    JsonMergeStrategy,
    MergeResult,
    SequentialMergeStrategy,
)
from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox
from cleveragents.infrastructure.sandbox.protocol import (
    CommitResult,
    Sandbox,
    SandboxContext,
    SandboxError,
    SandboxRollbackError,
    SandboxStateError,
    SandboxStatus,
)


def _sandbox_status_transitions() -> None:
    """Integration test: SandboxStatus state machine transitions."""
    # Valid transitions
    assert SandboxStatus.can_transition(SandboxStatus.PENDING, SandboxStatus.CREATED)
    assert SandboxStatus.can_transition(SandboxStatus.CREATED, SandboxStatus.ACTIVE)
    assert SandboxStatus.can_transition(SandboxStatus.ACTIVE, SandboxStatus.COMMITTED)
    assert SandboxStatus.can_transition(SandboxStatus.ACTIVE, SandboxStatus.ROLLED_BACK)
    assert SandboxStatus.can_transition(SandboxStatus.ACTIVE, SandboxStatus.ERRORED)

    # Invalid transitions
    assert not SandboxStatus.can_transition(
        SandboxStatus.COMMITTED, SandboxStatus.ACTIVE
    )
    assert not SandboxStatus.can_transition(
        SandboxStatus.CLEANED_UP, SandboxStatus.ACTIVE
    )

    # assert_transition should raise on invalid
    try:
        SandboxStatus.assert_transition(SandboxStatus.COMMITTED, SandboxStatus.ACTIVE)
        print("FAIL: Expected SandboxStateError")
        return
    except SandboxStateError:
        pass

    print("sandbox-status-transitions-ok")


def _no_sandbox_lifecycle() -> None:
    """Integration test: NoSandbox full lifecycle.

    Covers: create -> get_path -> commit -> cleanup.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        sandbox = NoSandbox(resource_id="res-001", original_path=tmpdir)

        # Verify protocol compliance
        assert isinstance(sandbox, Sandbox)
        assert sandbox.status == SandboxStatus.PENDING

        # Create
        ctx = sandbox.create(plan_id="plan-001")
        assert isinstance(ctx, SandboxContext)
        assert ctx.sandbox_path == tmpdir
        assert ctx.original_path == tmpdir
        assert sandbox.status == SandboxStatus.CREATED

        # Get path
        path = sandbox.get_path("src/main.py")
        assert path.endswith("src/main.py")
        assert sandbox.status == SandboxStatus.ACTIVE

        # Path traversal guard
        try:
            sandbox.get_path("../../../etc/passwd")
            print("FAIL: Expected SandboxError on path traversal")
            return
        except (SandboxError, ValueError):
            pass

        # Commit (no-op for NoSandbox)
        result = sandbox.commit("test commit")
        assert isinstance(result, CommitResult)
        assert result.success
        assert sandbox.status == SandboxStatus.COMMITTED

        # Cleanup
        sandbox.cleanup()
        assert sandbox.status == SandboxStatus.CLEANED_UP

    print("no-sandbox-lifecycle-ok")


def _no_sandbox_rollback_raises() -> None:
    """Integration test: NoSandbox rollback always raises."""
    with tempfile.TemporaryDirectory() as tmpdir:
        sandbox = NoSandbox(resource_id="res-002", original_path=tmpdir)
        sandbox.create(plan_id="plan-002")
        sandbox.get_path("file.txt")

        try:
            sandbox.rollback()
            print("FAIL: Expected SandboxRollbackError")
            return
        except SandboxRollbackError:
            pass

    print("no-sandbox-rollback-ok")


def _factory_create_none_strategy() -> None:
    """Integration test: SandboxFactory creates correct sandbox for each strategy."""
    factory = SandboxFactory()

    # 'none' strategy -> NoSandbox
    sandbox = factory.create_sandbox(
        resource_id="res-003",
        original_path="/tmp/test",
        sandbox_strategy="none",
    )
    assert isinstance(sandbox, NoSandbox)

    # 'git_worktree' strategy -> GitWorktreeSandbox (implemented in B4)
    from cleveragents.infrastructure.sandbox.git_worktree import GitWorktreeSandbox

    gwt_sandbox = factory.create_sandbox(
        resource_id="res-004",
        original_path="/tmp/test",
        sandbox_strategy="git_worktree",
    )
    assert isinstance(gwt_sandbox, GitWorktreeSandbox)

    # 'copy_on_write' strategy -> CopyOnWriteSandbox (implemented in B4)
    from cleveragents.infrastructure.sandbox.copy_on_write import CopyOnWriteSandbox

    cow_sandbox = factory.create_sandbox(
        resource_id="res-005",
        original_path="/tmp/test",
        sandbox_strategy="copy_on_write",
    )
    assert isinstance(cow_sandbox, CopyOnWriteSandbox)

    # Verify transaction_rollback now produces a TransactionSandbox
    from cleveragents.infrastructure.sandbox.transaction_sandbox import (
        TransactionSandbox,
    )

    tx_sandbox = factory.create_sandbox(
        resource_id="res-006",
        original_path=":memory:",
        sandbox_strategy="transaction_rollback",
    )
    assert isinstance(tx_sandbox, TransactionSandbox)

    print("factory-create-none-ok")


def _factory_supported_strategies() -> None:
    """Integration test: SandboxFactory strategy support checks."""
    # All three concrete implementations are now supported
    assert SandboxFactory.is_supported("none")
    assert SandboxFactory.is_supported("git_worktree")
    assert SandboxFactory.is_supported("copy_on_write")

    # transaction_rollback is now supported
    assert SandboxFactory.is_supported("transaction_rollback")

    # Unimplemented strategies are not supported
    assert not SandboxFactory.is_supported("snapshot")

    # Resource type strategy mapping (spec-aligned)
    git_checkout_strats = SandboxFactory.get_supported_strategies("git-checkout")
    assert "none" in git_checkout_strats
    assert "git_worktree" in git_checkout_strats
    assert "copy_on_write" in git_checkout_strats

    api_strats = SandboxFactory.get_supported_strategies("api_endpoint")
    assert "none" in api_strats

    # Unknown resource types default to ["none"]
    unknown_strats = SandboxFactory.get_supported_strategies("unknown_type")
    assert unknown_strats == ["none"]

    print("factory-supported-ok")


def _manager_lifecycle() -> None:
    """Integration test: SandboxManager get_or_create, commit_all, cleanup_all."""
    factory = SandboxFactory()
    manager = SandboxManager(factory=factory, cleanup_on_exit=False)

    with tempfile.TemporaryDirectory() as tmpdir:
        # Get or create a sandbox
        sandbox = manager.get_or_create_sandbox(
            plan_id="plan-010",
            resource_id="res-010",
            original_path=tmpdir,
            sandbox_strategy="none",
        )
        assert isinstance(sandbox, NoSandbox)
        assert sandbox.status == SandboxStatus.CREATED

        # Get same sandbox again (should return existing)
        same = manager.get_or_create_sandbox(
            plan_id="plan-010",
            resource_id="res-010",
            original_path=tmpdir,
            sandbox_strategy="none",
        )
        assert same.sandbox_id == sandbox.sandbox_id

        # List sandboxes
        sandboxes = manager.list_sandboxes("plan-010")
        assert len(sandboxes) == 1

        # Get single sandbox
        found = manager.get_sandbox("plan-010", "res-010")
        assert found is not None
        assert found.sandbox_id == sandbox.sandbox_id

        # Commit all
        results = manager.commit_all("plan-010")
        assert len(results) == 1
        assert results[0].success

        # Cleanup all
        manager.cleanup_all("plan-010")
        remaining = manager.list_sandboxes("plan-010")
        assert len(remaining) == 0

    print("manager-lifecycle-ok")


def _sequential_merge_strategy() -> None:
    """Integration test: SequentialMergeStrategy always returns theirs."""
    strategy = SequentialMergeStrategy()
    result = strategy.merge(
        base="original content",
        ours="our changes",
        theirs="their changes",
    )
    assert isinstance(result, MergeResult)
    assert result.success
    assert result.content == "their changes"
    assert not result.has_conflicts

    print("sequential-merge-ok")


def _json_merge_strategy() -> None:
    """Integration test: JsonMergeStrategy deep merges JSON."""
    strategy = JsonMergeStrategy(array_mode="replace")

    base = '{"a": 1, "b": {"c": 2}}'
    ours = '{"a": 1, "b": {"c": 3, "d": 4}}'
    theirs = '{"a": 10, "b": {"c": 5, "e": 6}}'

    result = strategy.merge(base=base, ours=ours, theirs=theirs)
    assert result.success
    merged = json.loads(result.content)
    # theirs overrides scalar values, but deep merge should combine nested keys
    assert merged["a"] == 10  # theirs wins
    assert "c" in merged["b"]
    assert "e" in merged["b"]

    print("json-merge-ok")


def _json_merge_concat_arrays() -> None:
    """Integration test: JsonMergeStrategy with concat array mode."""
    strategy = JsonMergeStrategy(array_mode="concat")

    base = '{"items": [1, 2]}'
    ours = '{"items": [1, 2, 3]}'
    theirs = '{"items": [4, 5]}'

    result = strategy.merge(base=base, ours=ours, theirs=theirs)
    assert result.success
    merged = json.loads(result.content)
    # concat mode should combine arrays
    assert len(merged["items"]) > 2

    print("json-merge-concat-ok")


def _manager_multiple_resources() -> None:
    """Integration test: SandboxManager handles multiple resources per plan."""
    factory = SandboxFactory()
    manager = SandboxManager(factory=factory, cleanup_on_exit=False)

    with (
        tempfile.TemporaryDirectory() as tmpdir1,
        tempfile.TemporaryDirectory() as tmpdir2,
    ):
        manager.get_or_create_sandbox(
            plan_id="plan-multi",
            resource_id="res-A",
            original_path=tmpdir1,
            sandbox_strategy="none",
        )
        manager.get_or_create_sandbox(
            plan_id="plan-multi",
            resource_id="res-B",
            original_path=tmpdir2,
            sandbox_strategy="none",
        )

        sandboxes = manager.list_sandboxes("plan-multi")
        assert len(sandboxes) == 2

        results = manager.commit_all("plan-multi")
        assert len(results) == 2
        assert all(r.success for r in results)

        manager.cleanup_all("plan-multi")
        assert len(manager.list_sandboxes("plan-multi")) == 0

    print("manager-multiple-resources-ok")


def _manager_atomic_commit() -> None:
    """Integration test: SandboxManager.commit_all is atomic.

    Verifies that when one sandbox commit fails, the entire batch fails
    and already-committed sandboxes are rolled back.
    """
    from datetime import datetime
    from unittest.mock import MagicMock, PropertyMock

    factory = SandboxFactory()
    manager = SandboxManager(factory=factory, cleanup_on_exit=False)

    with (
        tempfile.TemporaryDirectory() as tmpdir1,
        tempfile.TemporaryDirectory() as tmpdir2,
    ):
        # Create two sandboxes
        manager.get_or_create_sandbox(
            plan_id="plan-atomic",
            resource_id="res-ok",
            original_path=tmpdir1,
            sandbox_strategy="none",
        )
        manager.get_or_create_sandbox(
            plan_id="plan-atomic",
            resource_id="res-fail",
            original_path=tmpdir2,
            sandbox_strategy="none",
        )

        # Replace res-fail with a mock that fails on commit
        fail_mock = MagicMock()
        fail_mock.sandbox_id = "sb-fail"
        type(fail_mock).status = PropertyMock(return_value=SandboxStatus.ACTIVE)
        fail_mock.commit.side_effect = SandboxError("simulated commit failure")
        manager._active_sandboxes["plan-atomic"]["res-fail"] = fail_mock

        # Replace res-ok with a mock that commits and can be rolled back
        from cleveragents.infrastructure.sandbox.protocol import CommitResult

        ok_mock = MagicMock()
        ok_mock.sandbox_id = "sb-ok"
        type(ok_mock).status = PropertyMock(return_value=SandboxStatus.ACTIVE)
        ok_mock.commit.return_value = CommitResult(
            sandbox_id="sb-ok",
            success=True,
            timestamp=datetime.now(),
        )
        ok_mock.rollback.return_value = None
        manager._active_sandboxes["plan-atomic"]["res-ok"] = ok_mock

        results = manager.commit_all("plan-atomic")

        # Should return exactly one failure result
        assert len(results) == 1, f"Expected 1 result, got {len(results)}"
        assert results[0].success is False
        assert "Atomic commit failed" in (results[0].error or "")

        # The OK sandbox should have been rolled back
        assert ok_mock.rollback.called, "Expected rollback on already-committed sandbox"

        # Metadata should report rolled-back and failed-rollback IDs
        metadata = results[0].metadata
        assert "rolled_back" in metadata, "Expected 'rolled_back' in metadata"
        assert "rollback_failed" in metadata, "Expected 'rollback_failed' in metadata"

    print("manager-atomic-commit-ok")


def main() -> None:
    if len(sys.argv) < 2:
        raise SystemExit("Expected command argument")
    command = sys.argv[1]
    commands = {
        "status-transitions": _sandbox_status_transitions,
        "no-sandbox-lifecycle": _no_sandbox_lifecycle,
        "no-sandbox-rollback": _no_sandbox_rollback_raises,
        "factory-create-none": _factory_create_none_strategy,
        "factory-supported": _factory_supported_strategies,
        "manager-lifecycle": _manager_lifecycle,
        "sequential-merge": _sequential_merge_strategy,
        "json-merge": _json_merge_strategy,
        "json-merge-concat": _json_merge_concat_arrays,
        "manager-multiple-resources": _manager_multiple_resources,
        "manager-atomic-commit": _manager_atomic_commit,
    }
    if command not in commands:
        raise SystemExit(f"Unknown command: {command}")
    commands[command]()


if __name__ == "__main__":
    main()