import pytest

from app.core.errors import AgentError, RecoverableAgentError
from app.worker.agent_run_worker import AgentRunWorker


@pytest.mark.asyncio
async def test_worker_retries_recoverable_error() -> None:
    class FakeSession:
        def begin(self):
            class Context:
                async def __aenter__(self): return self
                async def __aexit__(self, *args): return False
            return Context()

    class FakeRepository:
        session = FakeSession()
        async def claim(self, run_id, worker_id, lease_seconds): return True
        async def record_failure(self, run_id, error_code, retry_limit): return True

    worker = object.__new__(AgentRunWorker)
    worker.repository = FakeRepository()
    worker.worker_id = "worker-1"
    worker.lease_seconds = 60
    worker.retry_limit = 3

    async def execute(_: str) -> None: raise RecoverableAgentError("temporary")
    worker.execute_run = execute
    assert await worker.claim_and_execute("run-1") is False


@pytest.mark.asyncio
async def test_worker_stops_non_recoverable_error() -> None:
    class FakeSession:
        def begin(self):
            class Context:
                async def __aenter__(self): return self
                async def __aexit__(self, *args): return False
            return Context()

    class FakeRepository:
        session = FakeSession()
        async def claim(self, run_id, worker_id, lease_seconds): return True
        async def record_failure(self, run_id, error_code, retry_limit): return False

    worker = object.__new__(AgentRunWorker)
    worker.repository = FakeRepository()
    worker.worker_id = "worker-1"
    worker.lease_seconds = 60
    worker.retry_limit = 3

    async def execute(_: str) -> None: raise AgentError("fatal")
    worker.execute_run = execute
    assert await worker.claim_and_execute("run-1") is False