dataloader/tests/integrations/test_worker_protocol.py

# tests/integrations/test_worker_protocol.py
import asyncio
from datetime import datetime, timedelta, timezone
from uuid import uuid4

import pytest
from sqlalchemy.ext.asyncio import AsyncSession

from dataloader.storage.repositories import QueueRepository, CreateJobRequest
from dataloader.context import APP_CTX

@pytest.mark.anyio
async def test_e2e_worker_protocol_ok(db_session: AsyncSession):
    """
    Проверяет полный E2E-сценарий жизненного цикла задачи:
    1. Постановка (create)
    2. Захват (claim)
    3. Пульс (heartbeat)
    4. Успешное завершение (finish_ok)
    5. Проверка статуса
    """
    repo = QueueRepository(db_session)
    job_id = str(uuid4())
    queue_name = "e2e_ok_queue"
    lock_key = f"lock_{job_id}"

    # 1. Постановка задачи
    create_req = CreateJobRequest(
        job_id=job_id,
        queue=queue_name,
        task="test_e2e_task",
        args={},
        idempotency_key=None,
        lock_key=lock_key,
        partition_key="",
        priority=100,
        available_at=datetime.now(timezone.utc),
        max_attempts=3,
        lease_ttl_sec=30,
        producer=None,
        consumer_group=None,
    )
    await repo.create_or_get(create_req)

    # 2. Захват задачи
    claimed_job = await repo.claim_one(queue_name, claim_backoff_sec=10)
    assert claimed_job is not None
    assert claimed_job["job_id"] == job_id
    assert claimed_job["lock_key"] == lock_key

    # 3. Пульс
    success, cancel_requested = await repo.heartbeat(job_id, ttl_sec=60)
    assert success
    assert not cancel_requested

    # 4. Успешное завершение
    await repo.finish_ok(job_id)

    # 5. Проверка статуса
    status = await repo.get_status(job_id)
    assert status is not None
    assert status.status == "succeeded"
    assert status.finished_at is not None


@pytest.mark.anyio
async def test_concurrency_claim_one_locks(db_session: AsyncSession):
    """
    Проверяет, что при конкурентном доступе к задачам с одинаковым
    lock_key только один воркер может захватить задачу.
    """
    repo = QueueRepository(db_session)
    queue_name = "concurrency_queue"
    lock_key = "concurrent_lock_123"
    job_ids = [str(uuid4()), str(uuid4())]

    # 1. Создание двух задач с одинаковым lock_key
    for i, job_id in enumerate(job_ids):
        create_req = CreateJobRequest(
            job_id=job_id,
            queue=queue_name,
            task=f"task_{i}",
            args={},
            idempotency_key=f"idem_con_{i}",
            lock_key=lock_key,
            partition_key="",
            priority=100 + i,
            available_at=datetime.now(timezone.utc),
            max_attempts=1,
            lease_ttl_sec=30,
            producer="test",
            consumer_group="test_group",
        )
        await repo.create_or_get(create_req)

    # 2. Первый воркер захватывает задачу
    claimed_job_1 = await repo.claim_one(queue_name, claim_backoff_sec=1)
    assert claimed_job_1 is not None
    assert claimed_job_1["job_id"] == job_ids[0]

    # 3. Второй воркер пытается захватить задачу, но не может (из-за advisory lock)
    claimed_job_2 = await repo.claim_one(queue_name, claim_backoff_sec=1)
    assert claimed_job_2 is None

    # 4. Первый воркер освобождает advisory lock (как будто завершил работу)
    await repo._advisory_unlock(lock_key)

    # 5. Второй воркер теперь может захватить вторую задачу
    claimed_job_3 = await repo.claim_one(queue_name, claim_backoff_sec=1)
    assert claimed_job_3 is not None
    assert claimed_job_3["job_id"] == job_ids[1]


@pytest.mark.anyio
async def test_reaper_requeues_lost_jobs(db_session: AsyncSession):
    """
    Проверяет, что reaper корректно возвращает "потерянные" задачи в очередь.
    """
    repo = QueueRepository(db_session)
    job_id = str(uuid4())
    queue_name = "reaper_queue"

    # 1. Создаем и захватываем задачу
    create_req = CreateJobRequest(
        job_id=job_id,
        queue=queue_name,
        task="reaper_test_task",
        args={},
        idempotency_key="idem_reaper_1",
        lock_key="reaper_lock_1",
        partition_key="",
        priority=100,
        available_at=datetime.now(timezone.utc),
        max_attempts=3,
        lease_ttl_sec=1,  # Очень короткий lease
        producer=None,
        consumer_group=None,
    )
    await repo.create_or_get(create_req)

    claimed_job = await repo.claim_one(queue_name, claim_backoff_sec=1)
    assert claimed_job is not None
    assert claimed_job["job_id"] == job_id

    # 2. Ждем истечения lease
    await asyncio.sleep(2)

    # 3. Запускаем reaper
    requeued_ids = await repo.requeue_lost()
    assert requeued_ids == [job_id]

    # 4. Проверяем статус
    status = await repo.get_status(job_id)
    assert status is not None
    assert status.status == "queued"