55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
|
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
|