hyungi_document_server/tests/test_worker_jobs_retry.py

"""PR-Worker-Pool-Registry-1B 정정 #3 — explicit /result failed 재시도 정책.

3 case:
  A. initial attempts=0, max_attempts=3 → claim 후 attempts=1 → failed → status='pending' 복귀.
     (worker_id=NULL, claimed_at=NULL, completed_at=NULL, error_message 저장)
  B. initial attempts=2, max_attempts=3 → claim 후 attempts=3 → failed → final 'failed'.
     (completed_at=now)
  C. failed 요청에 result={"partial":"..."} 포함 → DB result IS NULL 유지 (request.result 무시).
"""

from __future__ import annotations

import os
import sys
import uuid

import pytest
import pytest_asyncio

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "app"))

from httpx import ASGITransport, AsyncClient
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from _worker_pool_helpers import (
    cleanup_worker_capabilities,
    cleanup_worker_jobs,
    ensure_user,
    get_database_url,
    mint_access_token,
)


@pytest_asyncio.fixture
async def db_session(monkeypatch):
    monkeypatch.setenv("LAPTOP_WORKER_BOT_USERNAME", "laptop-worker-bot")
    engine = create_async_engine(get_database_url())
    sm = async_sessionmaker(engine, expire_on_commit=False)
    async with sm() as session:
        yield session
    await engine.dispose()


@pytest_asyncio.fixture
async def worker_token(db_session):
    await ensure_user(db_session, "laptop-worker-bot")
    return mint_access_token("laptop-worker-bot")


@pytest_asyncio.fixture
async def owner_id(db_session):
    return await ensure_user(db_session, "test-owner-retry-1b")


@pytest_asyncio.fixture
async def worker_id_unique(db_session, owner_id):
    wid = f"test-retry-1b-{uuid.uuid4().hex[:8]}"
    await db_session.execute(
        text(
            "INSERT INTO worker_capabilities (worker_id, user_id, device_label, "
            "worker_class, tier) VALUES (:w, :u, 'lbl', 'laptop', 'laptop_small')"
        ),
        {"w": wid, "u": owner_id},
    )
    await db_session.commit()
    yield wid
    await cleanup_worker_jobs(db_session, "test-retry-1b")
    await cleanup_worker_capabilities(db_session, wid)


async def _seed_processing_job(
    db_session, owner_id: int, worker_id: str, attempts: int, max_attempts: int = 3
) -> int:
    job_id = (
        await db_session.execute(
            text(
                "INSERT INTO worker_jobs (user_id, job_type, status, worker_id, "
                "attempts, max_attempts, claimed_at) "
                "VALUES (:u, :j, 'processing', :w, :a, :m, NOW()) RETURNING id"
            ),
            {
                "u": owner_id,
                "j": f"test-retry-1b-{uuid.uuid4().hex[:8]}",
                "w": worker_id,
                "a": attempts,
                "m": max_attempts,
            },
        )
    ).scalar_one()
    await db_session.commit()
    return job_id


@pytest.mark.asyncio
async def test_case_a_failed_then_pending(
    db_session, worker_token, worker_id_unique, owner_id
):
    """Case A: attempts(=1) < max_attempts(=3) → pending 복귀."""
    from main import app

    job_id = await _seed_processing_job(db_session, owner_id, worker_id_unique, attempts=1)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
        r = await c.post(
            "/internal/worker/result",
            json={
                "job_id": job_id,
                "worker_id": worker_id_unique,
                "status": "failed",
                "error_message": "case A error",
            },
            headers={"Authorization": f"Bearer {worker_token}"},
        )
    assert r.status_code == 200, r.text

    res = await db_session.execute(
        text(
            "SELECT status, worker_id, claimed_at, completed_at, error_message "
            "FROM worker_jobs WHERE id = :i"
        ),
        {"i": job_id},
    )
    s, w_id, c_at, comp_at, em = res.first()
    assert s == "pending"
    assert w_id is None
    assert c_at is None
    assert comp_at is None
    assert em == "case A error"


@pytest.mark.asyncio
async def test_case_b_failed_max_attempts_final(
    db_session, worker_token, worker_id_unique, owner_id
):
    """Case B: attempts(=3) >= max_attempts(=3) → final 'failed' + completed_at."""
    from main import app

    job_id = await _seed_processing_job(db_session, owner_id, worker_id_unique, attempts=3)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
        r = await c.post(
            "/internal/worker/result",
            json={
                "job_id": job_id,
                "worker_id": worker_id_unique,
                "status": "failed",
                "error_message": "case B final",
            },
            headers={"Authorization": f"Bearer {worker_token}"},
        )
    assert r.status_code == 200

    res = await db_session.execute(
        text("SELECT status, completed_at FROM worker_jobs WHERE id = :i"),
        {"i": job_id},
    )
    s, comp_at = res.first()
    assert s == "failed"
    assert comp_at is not None


@pytest.mark.asyncio
async def test_case_c_failed_ignores_result_field(
    db_session, worker_token, worker_id_unique, owner_id
):
    """Case C: failed 요청 result={'partial':...} 포함해도 DB result IS NULL 유지."""
    from main import app

    job_id = await _seed_processing_job(db_session, owner_id, worker_id_unique, attempts=1)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
        r = await c.post(
            "/internal/worker/result",
            json={
                "job_id": job_id,
                "worker_id": worker_id_unique,
                "status": "failed",
                "result": {"partial": "should be ignored"},
                "error_message": "case C",
            },
            headers={"Authorization": f"Bearer {worker_token}"},
        )
    assert r.status_code == 200
    res = await db_session.execute(
        text("SELECT result FROM worker_jobs WHERE id = :i"), {"i": job_id}
    )
    assert res.scalar_one() is None  # 정정 #3 엄격 규칙