hyungi_document_server/tests/test_worker_jobs_skip_locked.py

"""PR-Worker-Pool-Registry-1B — /claim 동시성 (SKIP LOCKED) + 204 body 검증.

2 항목:
  1. 두 async session 동시 claim → 한쪽만 200, 다른 쪽 204 + body 0 (정정 #4)
  2. queue empty + 별도 호출 → 204 + Content-Length: 0 명시
"""

from __future__ import annotations

import asyncio
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-skip-1b")


@pytest.mark.asyncio
async def test_skip_locked_only_one_winner(db_session, worker_token, owner_id):
    """두 client 가 동시에 동일 job_type claim → 한쪽 200, 다른 쪽 204."""
    from main import app

    jt = f"test-skip-1b-{uuid.uuid4().hex[:8]}"
    w1 = f"test-skip-1b-w1-{uuid.uuid4().hex[:6]}"
    w2 = f"test-skip-1b-w2-{uuid.uuid4().hex[:6]}"

    # seed capability + 1 pending job
    for w in (w1, w2):
        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') "
                "ON CONFLICT (worker_id) DO NOTHING"
            ),
            {"w": w, "u": owner_id},
        )
    await db_session.execute(
        text("INSERT INTO worker_jobs (user_id, job_type) VALUES (:u, :j)"),
        {"u": owner_id, "j": jt},
    )
    await db_session.commit()

    headers = {"Authorization": f"Bearer {worker_token}"}
    try:
        async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c1, \
                   AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c2:
            r1, r2 = await asyncio.gather(
                c1.post(
                    "/internal/worker/claim",
                    json={"worker_id": w1, "job_type": jt},
                    headers=headers,
                ),
                c2.post(
                    "/internal/worker/claim",
                    json={"worker_id": w2, "job_type": jt},
                    headers=headers,
                ),
            )
        codes = sorted([r1.status_code, r2.status_code])
        assert codes == [200, 204], f"unexpected codes: {codes}"
        loser = r1 if r1.status_code == 204 else r2
        # 정정 #4: 204 body 빈 검증
        assert loser.content == b""
    finally:
        await cleanup_worker_jobs(db_session, "test-skip-1b")
        await cleanup_worker_capabilities(db_session, "test-skip-1b-w")


@pytest.mark.asyncio
async def test_claim_204_body_explicit_empty(db_session, worker_token):
    """queue empty 호출 → 204 + content empty (정정 #4 명시적)."""
    from main import app

    async with AsyncClient(
        transport=ASGITransport(app=app), base_url="http://test"
    ) as c:
        r = await c.post(
            "/internal/worker/claim",
            json={"worker_id": "test-skip-1b-empty", "job_type": "test-skip-1b-NOEXIST"},
            headers={"Authorization": f"Bearer {worker_token}"},
        )
    assert r.status_code == 204
    assert r.content == b""
    # Content-Length 헤더는 ASGI 가 자동 0 으로 세팅. 없을 수도 있으므로 content 만 보장.