feat(worker-pool): Registry-1B Pull 활성화 (auth + worker_jobs + 5 endpoint)

worker-pool-policy §B 1B 영역 완료. 1A scaffold (mig 270~274 + 503 stub) 위에: - mig 275/276: worker_jobs (status CHECK + user_id=owner) + pending partial index - create_laptop_worker_bot_token + require_worker_user dependency (voice-memo 동형) - /internal/worker/{register,heartbeat,claim,result,drain} 5 endpoint 실 구현 - /claim FOR UPDATE SKIP LOCKED + 204 body 0 - /result 소유권 검증 (worker_id 매칭, 404) + failed 재시도 (attempts/max) - explicit failure 시 request.result 무시 (DB result NULL 유지) - 테스트 22 항목 7 파일 policy §B.2 5 invariant 보존: voice-memo wrapper 변경 0, drain advisory, result raw JSONB, ProcessingQueue 무변경, 운영 자동 분기 변경 0. 활용처 (recap context + /jobs/recap + payload 100KB guard) = Registry-1C 영역. stale recovery / 노트북 client / canonical promote = Notebook-Pilot-1 영역. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-19 08:53:57 +09:00
parent acd29b963e
commit f60d6e52fc
17 changed files with 1562 additions and 64 deletions
@@ -0,0 +1,275 @@
+"""PR-Worker-Pool-Registry-1B — /internal/worker/* 5 endpoint 정상 동작 검증.
+
+5 항목:
+  1. /register UPSERT — 같은 worker_id 두 번 호출하면 단일 row 유지 + last_registered_at 갱신.
+  2. /heartbeat — INSERT row 추가 + status='available' 저장 확인.
+  3. /claim — queue empty 시 204 + body 0. pending row 있으면 200 + payload + status='processing'.
+  4. /result — completed 시 status='completed' + result JSONB 저장.
+  5. /drain — heartbeat row status='draining' INSERT.
+
+DB integration. GPU docker compose exec fastapi pytest 환경 가정.
+"""
+
+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 client():
+    from main import app
+
+    async with AsyncClient(
+        transport=ASGITransport(app=app),
+        base_url="http://test",
+    ) as ac:
+        yield ac
+
+
+@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-user-1b")
+
+
+@pytest_asyncio.fixture
+async def worker_id_unique():
+    wid = f"test-mbp-1b-{uuid.uuid4().hex[:8]}"
+    yield wid
+
+
+@pytest_asyncio.fixture
+async def cleanup_after(db_session, worker_id_unique):
+    yield
+    await cleanup_worker_jobs(db_session, "test-rec-1b")
+    await cleanup_worker_capabilities(db_session, worker_id_unique)
+
+
+def _auth(token: str) -> dict:
+    return {"Authorization": f"Bearer {token}"}
+
+
+@pytest.mark.asyncio
+async def test_register_upsert(client, db_session, worker_token, worker_id_unique, cleanup_after):
+    body = {
+        "worker_id": worker_id_unique,
+        "device_label": "MacBook Pro M3 Pro",
+        "worker_class": "laptop",
+        "tier": "laptop_small",
+        "capabilities": ["short_summary"],
+        "models_loaded": ["gemma-3-4b-it"],
+    }
+    r1 = await client.post("/internal/worker/register", json=body, headers=_auth(worker_token))
+    assert r1.status_code == 200, r1.text
+    # 두번째 호출 = upsert 갱신
+    body["device_label"] = "MacBook Pro M3 Pro (renamed)"
+    r2 = await client.post("/internal/worker/register", json=body, headers=_auth(worker_token))
+    assert r2.status_code == 200
+    res = await db_session.execute(
+        text(
+            "SELECT device_label, last_registered_at FROM worker_capabilities WHERE worker_id = :w"
+        ),
+        {"w": worker_id_unique},
+    )
+    rows = res.all()
+    assert len(rows) == 1
+    assert rows[0][0] == "MacBook Pro M3 Pro (renamed)"
+
+
+@pytest.mark.asyncio
+async def test_heartbeat_insert(
+    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
+):
+    # 사전: worker_capabilities row 필요 (FK)
+    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": worker_id_unique, "u": owner_id},
+    )
+    await db_session.commit()
+
+    r = await client.post(
+        "/internal/worker/heartbeat",
+        json={"worker_id": worker_id_unique, "status": "available"},
+        headers=_auth(worker_token),
+    )
+    assert r.status_code == 200
+    res = await db_session.execute(
+        text(
+            "SELECT status FROM worker_heartbeats WHERE worker_id = :w ORDER BY heartbeat_at DESC LIMIT 1"
+        ),
+        {"w": worker_id_unique},
+    )
+    assert res.scalar_one() == "available"
+
+
+@pytest.mark.asyncio
+async def test_claim_204_when_empty(client, worker_token, cleanup_after):
+    r = await client.post(
+        "/internal/worker/claim",
+        json={"worker_id": "test-empty-1b", "job_type": "test-rec-1b-empty"},
+        headers=_auth(worker_token),
+    )
+    assert r.status_code == 204
+    # 정정 #4: body 0
+    assert r.content == b""
+
+
+@pytest.mark.asyncio
+async def test_claim_processing_transition(
+    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
+):
+    # seed worker_capabilities + worker_jobs
+    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": worker_id_unique, "u": owner_id},
+    )
+    jt = "test-rec-1b-claim"
+    job_id = (
+        await db_session.execute(
+            text(
+                "INSERT INTO worker_jobs (user_id, job_type, payload) "
+                "VALUES (:u, :j, '{\"week\":\"2026-W20\"}'::jsonb) RETURNING id"
+            ),
+            {"u": owner_id, "j": jt},
+        )
+    ).scalar_one()
+    await db_session.commit()
+
+    r = await client.post(
+        "/internal/worker/claim",
+        json={"worker_id": worker_id_unique, "job_type": jt},
+        headers=_auth(worker_token),
+    )
+    assert r.status_code == 200
+    js = r.json()
+    assert js["id"] == job_id
+    assert js["job_type"] == jt
+    assert js["attempts"] == 1
+    assert js["payload"]["week"] == "2026-W20"
+
+    res = await db_session.execute(
+        text("SELECT status, worker_id, attempts FROM worker_jobs WHERE id = :i"),
+        {"i": job_id},
+    )
+    status_, w_id, attempts = res.first()
+    assert status_ == "processing"
+    assert w_id == worker_id_unique
+    assert attempts == 1
+
+
+@pytest.mark.asyncio
+async def test_result_completed(
+    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
+):
+    # seed
+    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": worker_id_unique, "u": owner_id},
+    )
+    jt = "test-rec-1b-result"
+    job_id = (
+        await db_session.execute(
+            text(
+                "INSERT INTO worker_jobs (user_id, job_type, status, worker_id, "
+                "attempts, claimed_at) VALUES (:u, :j, 'processing', :w, 1, NOW()) RETURNING id"
+            ),
+            {"u": owner_id, "j": jt, "w": worker_id_unique},
+        )
+    ).scalar_one()
+    await db_session.commit()
+
+    r = await client.post(
+        "/internal/worker/result",
+        json={
+            "job_id": job_id,
+            "worker_id": worker_id_unique,
+            "status": "completed",
+            "result": {"summary": "test"},
+        },
+        headers=_auth(worker_token),
+    )
+    assert r.status_code == 200
+    res = await db_session.execute(
+        text("SELECT status, result FROM worker_jobs WHERE id = :i"), {"i": job_id}
+    )
+    s, result_jsonb = res.first()
+    assert s == "completed"
+    assert result_jsonb == {"summary": "test"}
+
+
+@pytest.mark.asyncio
+async def test_drain_heartbeat_insert(
+    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
+):
+    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": worker_id_unique, "u": owner_id},
+    )
+    await db_session.commit()
+
+    r = await client.post(
+        "/internal/worker/drain",
+        json={"worker_id": worker_id_unique, "reason": "sleep"},
+        headers=_auth(worker_token),
+    )
+    assert r.status_code == 200
+    res = await db_session.execute(
+        text(
+            "SELECT status, raw_payload FROM worker_heartbeats WHERE worker_id = :w "
+            "ORDER BY heartbeat_at DESC LIMIT 1"
+        ),
+        {"w": worker_id_unique},
+    )
+    s, rp = res.first()
+    assert s == "draining"
+    assert rp == {"reason": "sleep"}