Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e

tests/integration/jobs/test_job_engine.py

@@ -0,0 +1,185 @@
+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