RAGcore-Public/tests/integration/jobs/test_job_engine.py

from __future__ import annotations

import asyncio
import sys
from dataclasses import replace
from datetime import UTC, datetime, timedelta
from uuid import uuid4

import pytest
import pytest_asyncio
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from ragcore.application.jobs import (
    Job,
    JobContext,
    JobWorker,
    LeaseLostError,
    RetryableJobError,
    TerminalJobError,
)
from ragcore.config import Settings
from ragcore.infrastructure.database import create_database_engine, create_session_factory
from ragcore.infrastructure.database.jobs import PostgresJobRepository


@pytest.fixture(scope="session")
def event_loop_policy() -> asyncio.AbstractEventLoopPolicy:
    """psycopg's async mode refuses Windows' default proactor loop; the other
    PostgreSQL-backed modules already declare this, so this one does too."""

    if sys.platform == "win32":
        return asyncio.WindowsSelectorEventLoopPolicy()
    return asyncio.get_event_loop_policy()


async def _empty_the_queue(factory: async_sessionmaker[AsyncSession]) -> None:
    """Leave ``jobs.jobs`` as this test found it --- before *and* after.

    The queue is one shared table, and the ingestion and system gates later
    in a full run drive a real ``JobWorker``, which claims the oldest
    *available* job in it. A leftover fixture job is therefore claimed instead
    of theirs, and they fail for a reason that has nothing to do with what
    they test.
    """

    async with factory() as session:
        await session.execute(text("TRUNCATE jobs.events, jobs.attempts, jobs.jobs CASCADE"))
        await session.commit()


@pytest_asyncio.fixture
async def session_factory() -> async_sessionmaker[AsyncSession]:
    engine = create_database_engine(Settings(env="test"))
    factory = create_session_factory(engine)
    await _empty_the_queue(factory)
    try:
        yield factory
    finally:
        await _empty_the_queue(factory)
        await engine.dispose()


WORKSPACE_ID = uuid4()


def _job(
    now: datetime, *, key: str = "job-1", operation: str = "parse", max_attempts: int = 3
) -> Job:
    return Job(
        id=uuid4(),
        # Every job the platform enqueues is owned by exactly one workspace;
        # the repository refuses to write one that is not.
        workspace_id=WORKSPACE_ID,
        job_key=key,
        operation=operation,
        target_key="blob:abc",
        profile_version="parser-v1",
        payload={"source": "fixture"},
        available_at=now,
        max_attempts=max_attempts,
        trace_id="trace-job",
    )


@pytest.mark.integration
async def test_concurrent_workers_only_one_claims_a_ready_job(
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    now = datetime.now(UTC)
    first = PostgresJobRepository(session_factory)
    second = PostgresJobRepository(session_factory)
    await first.enqueue(_job(now))

    claims = await asyncio.gather(
        first.claim_next("worker-a", now=now, lease_duration=timedelta(minutes=1)),
        second.claim_next("worker-b", now=now, lease_duration=timedelta(minutes=1)),
    )

    assert sum(claim is not None for claim in claims) == 1


@pytest.mark.integration
async def test_expired_lease_is_reclaimed_and_old_owner_cannot_complete(
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    now = datetime.now(UTC)
    repository = PostgresJobRepository(session_factory)
    await repository.enqueue(_job(now))
    first = await repository.claim_next(
        "worker-a", now=now, lease_duration=timedelta(seconds=1)
    )
    assert first is not None

    reclaimed_at = now + timedelta(seconds=2)
    second = await repository.claim_next(
        "worker-b", now=reclaimed_at, lease_duration=timedelta(minutes=1)
    )
    assert second is not None
    assert second.attempt_number == 2
    with pytest.raises(LeaseLostError):
        await repository.complete(first, now=reclaimed_at)


@pytest.mark.integration
async def test_worker_distinguishes_retryable_and_terminal_errors(
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    now = datetime.now(UTC)
    repository = PostgresJobRepository(session_factory)
    retry_job = _job(now, key="retry", max_attempts=2)
    terminal_job = replace(
        _job(now, key="terminal", operation="terminal", max_attempts=2),
        available_at=now + timedelta(seconds=1),
    )
    await repository.enqueue(retry_job)
    await repository.enqueue(terminal_job)
    received: list[JobContext] = []

    async def retry(context: JobContext) -> None:
        received.append(context)
        raise RetryableJobError("temporary", retry_after_seconds=30)

    async def terminal(_context: JobContext) -> None:
        raise TerminalJobError("invalid input")

    worker = JobWorker(
        repository,
        {"parse": retry},
        worker_id="worker-retry",
        clock=lambda: now,
    )
    assert await worker.run_once()
    worker_terminal = JobWorker(
        repository,
        {"terminal": terminal},
        worker_id="worker-terminal",
        clock=lambda: now + timedelta(seconds=2),
    )
    assert await worker_terminal.run_once()

    assert received[0].idempotency_key == "retry"
    assert isinstance(received[0].cancellation, asyncio.Event)
    async with session_factory() as session:
        rows = (
            await session.execute(
                text("SELECT job_key, status, last_error_code FROM jobs.jobs ORDER BY job_key")
            )
        ).all()
    assert rows == [
        ("retry", "retry_wait", "retryable_error"),
        ("terminal", "terminal", "terminal_error"),
    ]
    async with session_factory() as session:
        event_rows = (
            await session.execute(
                text(
                    "SELECT event_type, error_code FROM jobs.events "
                    "WHERE event_type IN ('retry_scheduled', 'terminal') "
                    "ORDER BY event_type"
                )
            )
        ).all()
    assert ("retry_scheduled", "retryable_error") in event_rows
    assert ("terminal", "terminal_error") in event_rows