harness_zbt/backend/tests/test_refund_worker_reporting.py

from datetime import UTC, datetime
from threading import Event
from typing import Any

from zbt.commands.worker_reporting import RefundWorkerReporter
from zbt.domains.surrender.monitoring import InMemoryWorkerRegistry, WorkerInstance

NOW = datetime(2026, 7, 29, 11, 0, tzinfo=UTC)


def instance() -> WorkerInstance:
    return WorkerInstance(
        instance_id="worker-reporting-01",
        worker_name="refund",
        mode="parallel",
        host="node-a",
        pid=33001,
        status="RUNNING",
        started_at=NOW,
        last_heartbeat_at=NOW,
        last_poll_at=None,
        last_processed_at=None,
        total_processed=0,
        last_error=None,
    )


def test_reporter_records_worker_poll_and_graceful_stop() -> None:
    registry = InMemoryWorkerRegistry()
    reporter = RefundWorkerReporter(
        registry,
        instance(),
        clock=lambda: NOW,
        heartbeat_interval_seconds=60,
    )

    class SuccessfulWorker:
        def run_once(self, limit: int = 20) -> dict[str, Any]:
            assert limit == 20
            return {"processed": 2, "items": []}

    worker = reporter.wrap(SuccessfulWorker())
    result = worker.run_once()
    reporter.stop()
    reported = registry.list_instances("refund")[0]

    assert result["processed"] == 2
    assert reported.total_processed == 2
    assert reported.last_poll_at == NOW
    assert reported.status == "STOPPED"


def test_reporter_records_unhandled_worker_error() -> None:
    registry = InMemoryWorkerRegistry()
    reporter = RefundWorkerReporter(
        registry,
        instance(),
        clock=lambda: NOW,
        heartbeat_interval_seconds=60,
    )

    class FailedWorker:
        def run_once(self, limit: int = 20) -> dict[str, Any]:
            del limit
            raise RuntimeError("数据库连接中断")

    worker = reporter.wrap(FailedWorker())
    try:
        worker.run_once()
    except RuntimeError:
        pass
    else:
        raise AssertionError("工作器异常不应被监控层吞掉")
    reporter.stop(error="数据库连接中断")
    reported = registry.list_instances("refund")[0]

    assert reported.status == "ERROR"
    assert reported.last_error == "数据库连接中断"


def test_reporter_refreshes_heartbeat_while_worker_is_idle() -> None:
    registry = InMemoryWorkerRegistry()
    heartbeat_seen = Event()
    clock_calls = 0

    def clock() -> datetime:
        nonlocal clock_calls
        clock_calls += 1
        if clock_calls >= 1:
            heartbeat_seen.set()
        return NOW

    reporter = RefundWorkerReporter(
        registry,
        instance(),
        clock=clock,
        heartbeat_interval_seconds=0.001,
    )

    assert heartbeat_seen.wait(0.5)
    reporter.stop()
    assert clock_calls >= 1