RAG/backend/tests/unit/test_job_queue.py

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)