asf/services/asf-worker/tests/test_worker_handlers.py

"""Tests for worker handlers + claim/persistence logic.

All tests run in-memory using FakeJobsCollection — no Mongo/Rabbit required.
"""

from __future__ import annotations

from uuid import uuid4

import pytest

from asf_domain import JobStatus, JobUsage


class TestEchoHandler:
    @pytest.mark.asyncio
    async def test_echo_handler_returns_payload_as_result(self, worker_deps, sample_envelope):
        from app.handlers import echo_handler
        res = await echo_handler(sample_envelope, worker_deps)
        assert res.status == JobStatus.COMPLETED
        assert res.result["echo"] == sample_envelope.payload
        assert res.result["handled_by"] == worker_deps.worker_id
        assert isinstance(res.usage, JobUsage)
        assert res.usage.cpu_seconds > 0

    @pytest.mark.asyncio
    async def test_noop_handler_marks_failed_when_queue_unregistered(self, worker_deps, sample_envelope):
        from app.handlers import noop_handler
        res = await noop_handler(sample_envelope, worker_deps)
        assert res.status == JobStatus.FAILED
        assert "No worker handler registered" in res.error
        assert res.error_detail is not None and "hint" in res.error_detail


class TestMarkJobDonePersistence:
    @pytest.mark.asyncio
    async def test_mark_job_done_writes_status_result_usage(self, worker_deps, sample_envelope, fake_jobs_collection):
        from app.handlers import mark_job_done, HandlerResult
        fake_jobs_collection._docs[sample_envelope.job_id] = {
            "_id": sample_envelope.job_id,
            "status": JobStatus.RUNNING.value,
        }
        result = HandlerResult(
            status=JobStatus.COMPLETED,
            result={"echo": {"ok": True}},
            usage=JobUsage(cpu_seconds=0.042, input_tokens=100, output_tokens=200),
        )
        await mark_job_done(worker_deps, sample_envelope.job_id, result)
        doc = fake_jobs_collection._docs[sample_envelope.job_id]
        assert doc["status"] == JobStatus.COMPLETED.value
        assert doc["result"] == result.result
        assert doc["worker"] == worker_deps.worker_id
        assert doc["usage"]["cpu_seconds"] == 0.042
        assert doc["usage"]["input_tokens"] == 100
        assert doc["completed_at"] is not None

    @pytest.mark.asyncio
    async def test_mark_job_done_records_error_when_failed(self, worker_deps, sample_envelope, fake_jobs_collection):
        from app.handlers import mark_job_done, HandlerResult
        fake_jobs_collection._docs[sample_envelope.job_id] = {"_id": sample_envelope.job_id}
        await mark_job_done(
            worker_deps,
            sample_envelope.job_id,
            HandlerResult(
                status=JobStatus.FAILED,
                error="BOOM",
                error_detail={"stage": "analyze"},
            ),
        )
        doc = fake_jobs_collection._docs[sample_envelope.job_id]
        assert doc["status"] == JobStatus.FAILED.value
        assert doc["error"] == "BOOM"
        assert doc["error_detail"] == {"stage": "analyze"}
        # result / usage keys are ABSENT when empty
        assert "result" not in doc


class TestClaimRunningAtomic:
    @pytest.mark.asyncio
    async def test_claim_sets_running_and_inc_retries(self, worker_deps, sample_envelope, fake_jobs_collection):
        from app.worker import _claim_running
        fake_jobs_collection._docs[sample_envelope.job_id] = {
            "_id": sample_envelope.job_id,
            "tenant_id": sample_envelope.tenant_id,
            "status": JobStatus.QUEUED.value,
            "retries": 0,
        }
        await _claim_running(worker_deps, sample_envelope)
        doc = fake_jobs_collection._docs[sample_envelope.job_id]
        assert doc["status"] == JobStatus.RUNNING.value
        assert doc["worker"] == worker_deps.worker_id
        assert doc["retries"] == 1
        assert doc["started_at"] is not None

    @pytest.mark.asyncio
    async def test_claim_bails_if_job_already_terminal(self, worker_deps, sample_envelope, fake_jobs_collection):
        from app.worker import _claim_running
        fake_jobs_collection._docs[sample_envelope.job_id] = {
            "_id": sample_envelope.job_id,
            "tenant_id": sample_envelope.tenant_id,
            "status": JobStatus.COMPLETED.value,  # not QUEUED
            "worker": "some-other-worker",
        }
        await _claim_running(worker_deps, sample_envelope)
        doc = fake_jobs_collection._docs[sample_envelope.job_id]
        # still terminal + still original worker -> no overwrites
        assert doc["status"] == JobStatus.COMPLETED.value
        assert doc["worker"] == "some-other-worker"
        assert "retries" not in doc or doc["retries"] == 0


class TestAIMockedGatewayHandler:
    """Ensure ai_handler works safely with no real AI gateway attached (fallback path)."""

    @pytest.mark.asyncio
    async def test_ai_handler_without_gateway_returns_payload_note(self, worker_deps, sample_envelope):
        from asf_domain import JobType
        sample_envelope.type = JobType.AI_INTERACTIVE.value
        sample_envelope.payload = {"kind": "ask_asf", "question": "Summarize REQ-103"}
        worker_deps.ai_gateway = None
        from app.handlers import ai_handler
        res = await ai_handler(sample_envelope, worker_deps)
        assert res.status == JobStatus.COMPLETED
        assert "ai_gateway not attached" in res.result["note"]
        assert res.result["payload"]["kind"] == "ask_asf"


class TestHandlersQueueContract:
    def test_expected_queues_have_registered_handlers(self):
        from app.worker import HANDLERS_BY_QUEUE
        expected = {
            "asf.ai.interactive",
            "asf.ai.background",
            "asf.indexing",
            "asf.code_analysis",
            "asf.testing",
            "asf.devops",
            "asf.notifications",
        }
        assert set(HANDLERS_BY_QUEUE.keys()) == expected