microfish/backend/tests/test_job_queue.py

"""TDD gate: durable job queue claiming/consumption without a broker.

A real worker later runs on a queue provider, but the claim/complete lifecycle
and tenant scope must work against the durable ``jobs`` table now so jobs
survive restarts and are isolated per organization.
"""

import pytest
from sqlalchemy import create_engine

from app.db import Base, create_session_factory
from app.models.operations import Job, JobStatus
from app.services.job_queue import JobQueue


@pytest.fixture()
def session_factory():
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    factory = create_session_factory(engine)
    try:
        yield factory
    finally:
        engine.dispose()


def _seed_job(session_factory, *, organization_id="org-a", operation="graph.build"):
    session = session_factory()
    job = Job(
        organization_id=organization_id,
        owner_user_id="user-1",
        operation=operation,
        status=JobStatus.QUEUED,
    )
    session.add(job)
    session.commit()
    job_id = job.id
    session.close()
    return job_id


def test_job_queue_claims_single_queued_job(session_factory):
    job_id = _seed_job(session_factory)

    session = session_factory()
    try:
        queue = JobQueue(session)
        claimed = queue.claim_next_job(worker_id="worker-1")
        assert claimed is not None
        assert claimed.id == job_id
        assert claimed.status == JobStatus.RUNNING.value
    finally:
        session.close()


def test_job_queue_does_not_double_claim(session_factory):
    _seed_job(session_factory)

    s1 = session_factory()
    s2 = session_factory()
    try:
        q1 = JobQueue(s1)
        q2 = JobQueue(s2)
        first = q1.claim_next_job(worker_id="w-1")
        second = q2.claim_next_job(worker_id="w-2")
        assert first is not None
        # Second worker must not see the already-claimed job (and no other jobs).
        assert second is None
    finally:
        s1.close()
        s2.close()


def test_job_queue_does_not_claim_other_tenants_job(session_factory):
    _seed_job(session_factory, organization_id="org-a")
    _seed_job(session_factory, organization_id="org-b")

    session = session_factory()
    try:
        queue = JobQueue(session)
        # A worker processing org-a claims only org-a jobs.
        claimed = queue.claim_next_job(worker_id="w-1", organization_id="org-a")
        assert claimed is not None and claimed.organization_id == "org-a"
    finally:
        session.close()


def test_job_queue_complete_and_fail_update_status(session_factory):
    job_id = _seed_job(session_factory)

    session = session_factory()
    try:
        queue = JobQueue(session)
        claimed = queue.claim_next_job(worker_id="w-1")
        assert claimed is not None
        queue.complete_job(claimed.id, result={"ok": True})
        session.commit()
        refreshed = session.get(Job, claimed.id)
        assert refreshed.status == JobStatus.SUCCEEDED.value
        assert refreshed.result == {"ok": True}

        job_id2 = _seed_job(session_factory)
        claimed2 = queue.claim_next_job(worker_id="w-1")
        assert claimed2 is not None and claimed2.id == job_id2
        queue.fail_job(claimed2.id, error_code="boom")
        session.commit()
        refreshed2 = session.get(Job, claimed2.id)
        assert refreshed2.status == JobStatus.FAILED.value
        assert refreshed2.error_code == "boom"
    finally:
        session.close()


def test_job_queue_none_when_empty(session_factory):
    session = session_factory()
    try:
        queue = JobQueue(session)
        assert queue.claim_next_job(worker_id="w-1") is None
    finally:
        session.close()


def test_job_queue_dispatch_invokes_registered_handler(session_factory):
    calls = {}

    def handler(payload, job):
        calls["payload"] = payload
        calls["job_id"] = job.id
        return {"processed": True}

    job_id = _seed_job(session_factory, operation="graph.build")

    session = session_factory()
    try:
        queue = JobQueue(session)
        queue.register_handler("graph.build", handler)
        claimed = queue.claim_next_job(worker_id="w-1")
        assert claimed is not None
        result = queue.dispatch(claimed, payload={"graph_id": "g1"})
        assert result == {"processed": True}
        assert calls["job_id"] == job_id
        assert calls["payload"] == {"graph_id": "g1"}
    finally:
        session.close()


def test_job_queue_dispatch_fails_unhandled_operation(session_factory):
    job_id = _seed_job(session_factory, operation="unknown.op")

    session = session_factory()
    try:
        queue = JobQueue(session)
        claimed = queue.claim_next_job(worker_id="w-1")
        assert claimed is not None
        with pytest.raises(ValueError, match="no_handler"):
            queue.dispatch(claimed, payload={})
        # Claimed job is still running until the worker decides to fail it.
        refreshed = session.get(Job, claimed.id)
        assert refreshed.status == JobStatus.RUNNING.value
    finally:
        session.close()