cleveragents-core/features/steps/sandbox_manager_concurrency_steps.py

"""Concurrency-related step definitions for sandbox manager commit_all."""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass
from datetime import datetime
from typing import Any

from behave import given, then, when

from cleveragents.infrastructure.sandbox.manager import SandboxManager
from cleveragents.infrastructure.sandbox.protocol import CommitResult, SandboxStatus


@dataclass
class _ConcurrencyHarness:
    commit_started: threading.Event
    allow_commit: threading.Event
    commit_finished: threading.Event
    creation_done: threading.Event
    creation_error: Exception | None = None
    creation_completed_before_release: bool = False


class _BlockingSandbox:
    """Test sandbox that blocks commit until allowed."""

    def __init__(
        self,
        sandbox_id: str,
        commit_started: threading.Event,
        allow_commit: threading.Event,
    ) -> None:
        self._sandbox_id = sandbox_id
        self._status = SandboxStatus.CREATED
        self._commit_started = commit_started
        self._allow_commit = allow_commit

    @property
    def sandbox_id(self) -> str:
        return self._sandbox_id

    @property
    def status(self) -> SandboxStatus:
        return self._status

    def commit(self) -> CommitResult:
        self._commit_started.set()
        if not self._allow_commit.wait(timeout=1):
            raise TimeoutError("commit did not receive release signal in test harness")
        self._status = SandboxStatus.COMMITTED
        return CommitResult(
            sandbox_id=self._sandbox_id,
            success=True,
            timestamp=datetime.now(),
        )

    def rollback(self) -> None:  # pragma: no cover - not triggered in test
        self._status = SandboxStatus.ROLLED_BACK

    def cleanup(self) -> None:  # pragma: no cover - not used in test
        self._status = SandboxStatus.CLEANED_UP


@given('plan "{plan_id}" has a blocking sandbox commit harness')
def step_given_plan_has_blocking_sandbox(context: Any, plan_id: str) -> None:
    manager: SandboxManager = context.manager
    harness = _ConcurrencyHarness(
        commit_started=threading.Event(),
        allow_commit=threading.Event(),
        commit_finished=threading.Event(),
        creation_done=threading.Event(),
    )
    blocking = _BlockingSandbox(
        sandbox_id="sb-block",
        commit_started=harness.commit_started,
        allow_commit=harness.allow_commit,
    )

    with manager._lock:  # type: ignore[attr-defined]
        manager._active_sandboxes[plan_id] = {"res-block": blocking}

    context.concurrency_harness = harness
    context.plan_under_test = plan_id
    context.commit_results = None
    context.commit_error = None
    context.created_sandbox = None


@when(
    'commit_all runs concurrently with sandbox creation for plan "{plan_id}" resource "{resource_id}"'
)
def step_when_commit_all_runs_concurrently(
    context: Any, plan_id: str, resource_id: str
) -> None:
    manager: SandboxManager = context.manager
    harness: _ConcurrencyHarness = context.concurrency_harness

    def _run_commit() -> None:
        try:
            context.commit_results = manager.commit_all(plan_id)
        except Exception as exc:  # pragma: no cover - unexpected path
            context.commit_error = exc
        finally:
            harness.commit_finished.set()

    def _run_creation() -> None:
        try:
            context.created_sandbox = manager.get_or_create_sandbox(
                plan_id=plan_id,
                resource_id=resource_id,
                original_path="/tmp/resource",
                sandbox_strategy="none",
            )
        except Exception as exc:  # pragma: no cover - unexpected path
            harness.creation_error = exc
        finally:
            harness.creation_done.set()

    commit_thread = threading.Thread(target=_run_commit, daemon=True)
    creation_thread = threading.Thread(target=_run_creation, daemon=True)

    commit_thread.start()
    started = harness.commit_started.wait(timeout=1)
    assert started, "commit_all did not reach sandbox commit phase"

    creation_thread.start()
    time.sleep(0.05)
    harness.creation_completed_before_release = harness.creation_done.is_set()

    harness.allow_commit.set()

    commit_thread.join(timeout=1)
    creation_thread.join(timeout=1)

    assert harness.commit_finished.is_set(), "commit_all did not finish in test window"
    assert harness.creation_done.is_set(), "sandbox creation thread did not finish"
    assert context.commit_error is None, (
        f"Unexpected commit error: {context.commit_error}"
    )


@then("the concurrent sandbox creation should wait for commit completion")
def step_then_creation_waited(context: Any) -> None:
    harness: _ConcurrencyHarness = context.concurrency_harness
    assert harness.creation_completed_before_release is False, (
        "Sandbox creation should have been blocked until commit finished"
    )
    assert harness.creation_error is None, (
        f"Sandbox creation failed: {harness.creation_error}"
    )
    assert context.created_sandbox is not None, "Sandbox creation never completed"


@then("commit_all should return {expected:d} successful result")
def step_then_commit_results_success(context: Any, expected: int) -> None:
    results = context.commit_results
    assert results is not None, "commit_all did not return results"
    assert len(results) == expected, (
        f"Expected {expected} commit results, received {len(results)}"
    )
    assert all(result.success for result in results), "Not all commit results succeeded"