from __future__ import annotations import uuid from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock import pytest from app.persistence.job_queue import ( InvalidJobRowError, JobLease, LeaseLostError, PsycopgJobQueue, ) NOW = datetime(2026, 7, 13, 8, 0, tzinfo=UTC) JOB_ID = uuid.UUID("10000000-0000-0000-0000-000000000001") RESOURCE_ID = uuid.UUID("20000000-0000-0000-0000-000000000002") LEASE_TOKEN = uuid.UUID("30000000-0000-0000-0000-000000000003") def _claimed_row(*, worker_id: str = "worker-a") -> dict[str, object]: return { "id": JOB_ID, "job_type": "EMBED_DOCUMENT", "required_capability": "embedding", "resource_type": "document_version", "resource_id": RESOURCE_ID, "idempotency_key": "embed:version-1:profile-1", "payload": {"document_version_id": str(RESOURCE_ID)}, "stage": "EMBEDDING", "status": "RUNNING", "progress": 20, "priority": 5, "attempt": 1, "max_attempts": 3, "run_after": NOW, "lease_owner": worker_id, "lease_token": LEASE_TOKEN, "lease_until": NOW + timedelta(seconds=60), "created_at": NOW - timedelta(minutes=1), "updated_at": NOW, "finished_at": None, } def _state_row(*, status: str = "SUCCEEDED") -> dict[str, object]: return { "id": JOB_ID, "status": status, "attempt": 1, "max_attempts": 3, "finished_at": NOW if status in {"SUCCEEDED", "FAILED"} else None, } def _repository_with_rows( *, one: dict[str, object] | None = None, many: list[dict[str, object]] | None = None, ) -> tuple[PsycopgJobQueue, MagicMock, MagicMock, MagicMock]: cursor = MagicMock() cursor.fetchone.return_value = one cursor.fetchall.return_value = many or [] transaction = MagicMock() connection = MagicMock() connection.__enter__.return_value = connection connection.transaction.return_value = transaction connection.execute.return_value = cursor factory = MagicMock(return_value=connection) repository = PsycopgJobQueue( "postgresql://worker:private@db/rag", connection_factory=factory, ) return repository, factory, connection, transaction def test_claim_runs_in_short_transaction_and_returns_complete_fence() -> None: repository, factory, connection, transaction = _repository_with_rows(one=_claimed_row()) job = repository.claim( worker_id="worker-a", worker_capabilities=("embedding", "document_parse"), lease_seconds=60, ) assert job is not None assert job.id == JOB_ID assert job.payload == {"document_version_id": str(RESOURCE_ID)} assert job.lease == JobLease(JOB_ID, "worker-a", LEASE_TOKEN) factory.assert_called_once_with("postgresql://worker:private@db/rag", 5) transaction.__enter__.assert_called_once_with() transaction.__exit__.assert_called_once() statement, parameters = connection.execute.call_args.args assert "FOR UPDATE SKIP LOCKED" in statement assert "%(worker_id)s" in statement assert ":worker_id" not in statement assert parameters == { "worker_id": "worker-a", "worker_capabilities": ["embedding", "document_parse"], "lease_seconds": 60, } def test_claim_returns_none_without_fabricating_a_lease() -> None: repository, _, _, _ = _repository_with_rows(one=None) result = repository.claim( worker_id="worker-a", worker_capabilities=("embedding",), lease_seconds=60, ) assert result is None def test_claim_rejects_a_row_not_fenced_to_the_requesting_worker() -> None: repository, _, _, _ = _repository_with_rows(one=_claimed_row(worker_id="worker-b")) with pytest.raises(InvalidJobRowError, match="unexpected lease owner"): repository.claim( worker_id="worker-a", worker_capabilities=("embedding",), lease_seconds=60, ) def test_heartbeat_requires_owner_and_token_and_fails_closed() -> None: heartbeat_row = {"id": JOB_ID, "lease_until": NOW + timedelta(seconds=60)} repository, _, connection, _ = _repository_with_rows(one=heartbeat_row) lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN) heartbeat = repository.heartbeat(lease, lease_seconds=60) assert heartbeat.job_id == JOB_ID heartbeat_statement, parameters = connection.execute.call_args.args assert "job.lease_until >= now()" in heartbeat_statement assert parameters == { "job_id": JOB_ID, "worker_id": "worker-a", "lease_token": LEASE_TOKEN, "lease_seconds": 60, } lost_repository, _, _, _ = _repository_with_rows(one=None) with pytest.raises(LeaseLostError, match="no longer owned"): lost_repository.heartbeat(lease, lease_seconds=60) def test_complete_and_failure_updates_carry_the_full_fence() -> None: lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN) complete_repository, _, complete_connection, _ = _repository_with_rows(one=_state_row()) completed = complete_repository.complete(lease) assert completed.status == "SUCCEEDED" complete_statement, complete_parameters = complete_connection.execute.call_args.args assert "job.lease_until >= now()" in complete_statement assert complete_parameters == { "job_id": JOB_ID, "worker_id": "worker-a", "lease_token": LEASE_TOKEN, } retry_repository, _, retry_connection, _ = _repository_with_rows( one=_state_row(status="QUEUED") ) retried = retry_repository.fail_or_retry( lease, error_code="MODEL_TIMEOUT", error_message="Safe bounded failure", retry_delay_seconds=30, ) assert retried.status == "QUEUED" retry_statement, retry_parameters = retry_connection.execute.call_args.args assert "job.lease_until >= now()" in retry_statement assert retry_parameters == { "job_id": JOB_ID, "worker_id": "worker-a", "lease_token": LEASE_TOKEN, "error_code": "MODEL_TIMEOUT", "error_message": "Safe bounded failure", "retry_delay_seconds": 30, } def test_terminal_update_with_no_returning_row_is_lease_loss() -> None: repository, _, _, _ = _repository_with_rows(one=None) lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN) with pytest.raises(LeaseLostError): repository.complete(lease) with pytest.raises(LeaseLostError): repository.fail_or_retry( lease, error_code="JOB_HANDLER_FAILED", error_message="Safe failure", retry_delay_seconds=30, ) def test_reaper_uses_advisory_lock_and_bounded_batch() -> None: repository, _, connection, transaction = _repository_with_rows( many=[_state_row(status="QUEUED"), _state_row(status="FAILED")] ) states = repository.reap_expired(lock_key=42, batch_size=25) assert [state.status for state in states] == ["QUEUED", "FAILED"] statement, parameters = connection.execute.call_args.args assert "pg_try_advisory_xact_lock" in statement assert "FOR UPDATE OF job SKIP LOCKED" in statement assert parameters == {"lock_key": 42, "batch_size": 25} transaction.__exit__.assert_called_once() @pytest.mark.parametrize( ("operation", "expected_message"), [ ("empty_capability", "worker_capabilities"), ("duplicate_capability", "duplicates"), ("bad_error_code", "stable uppercase"), ("bad_batch", "batch_size"), ], ) def test_repository_rejects_invalid_operational_inputs( operation: str, expected_message: str, ) -> None: repository, _, _, _ = _repository_with_rows(one=None) lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN) with pytest.raises(ValueError, match=expected_message): if operation == "empty_capability": repository.claim( worker_id="worker-a", worker_capabilities=(), lease_seconds=60, ) elif operation == "duplicate_capability": repository.claim( worker_id="worker-a", worker_capabilities=("embedding", "embedding"), lease_seconds=60, ) elif operation == "bad_error_code": repository.fail_or_retry( lease, error_code="contains spaces", error_message="Safe failure", retry_delay_seconds=30, ) else: repository.reap_expired(lock_key=42, batch_size=0)