re-xxm/tests/test_task_state.py

# -*- coding: utf-8 -*-
"""task:state 持久化与启动续跑: uv run python tests/test_task_state.py"""

import json
import os
import sys
import tempfile
from datetime import datetime
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.stdout.reconfigure(encoding="utf-8")

from reaction_board.database import Database  # noqa: E402
from reaction_board.repository import Repository  # noqa: E402
from reaction_board.scheduler import ScheduledTask, Scheduler  # noqa: E402
from reaction_board.task_state import (  # noqa: E402
    STATUS_DONE,
    STATUS_FAILED,
    STATUS_PENDING,
    STATUS_RUNNING,
    ensure_pending,
    load_state,
    mark_done,
    mark_failed,
    mark_running,
    recover_stale_running,
    resume_stat_date,
    startup_next_run,
)
from reaction_board.timeutil import SHANGHAI_TZ, next_daily_run  # noqa: E402


def _repo(tmp: str) -> tuple[Database, Repository]:
    db = Database(Path(tmp) / "app.db")
    db.initialize()
    return db, Repository(db)


def test_restart_mid_run_resumes_same_stat_date():
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 23, 40, tzinfo=SHANGHAI_TZ).timestamp()
            mark_running(
                repo, "snapshot", stat_date="2026-11-03", now=now,
                progress={"done": 10, "total": 100},
            )
            # 模拟进程崩溃：running 残留
            recovered = recover_stale_running(repo, "snapshot", now=now + 60)
            assert recovered is not None
            assert recovered.status == STATUS_PENDING
            assert recovered.stat_date == "2026-11-03"
            assert recovered.progress == {"done": 10, "total": 100}
            assert recovered.due_at <= now + 60
            assert resume_stat_date(repo, "snapshot", fallback="2026-11-04") == "2026-11-03"
            nxt = startup_next_run(
                repo, "snapshot",
                wall_due_at=next_daily_run(now + 60, "23:15"),
                now=now + 60,
            )
            assert nxt <= now + 60
        finally:
            db.close()


def test_seeded_before_snapshot_at_does_not_force_run():
    """seed 后挂 pending due_at=当晚；白天重启 next_run 未到，不跑。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 10, 30, tzinfo=SHANGHAI_TZ).timestamp()
            slot = datetime(2026, 11, 3, 23, 15, tzinfo=SHANGHAI_TZ).timestamp()
            ensure_pending(
                repo, "snapshot", due_at=slot, stat_date="2026-11-03", now=now,
            )
            repo.kv_set("snapshot:seeded:2026-11-03", "1")
            nxt = startup_next_run(
                repo, "snapshot",
                wall_due_at=next_daily_run(now, "23:15"),
                now=now,
            )
            assert nxt == slot
            assert nxt > now
            calls = []

            def run_snapshot():
                calls.append("ran")
                return None

            clock = [now]
            sched = Scheduler(
                [ScheduledTask("snapshot", 86400.0, run_snapshot, next_run=nxt)],
                clock=lambda: clock[0],
            )
            assert sched.run_due() == []
            assert calls == []
        finally:
            db.close()


def test_failed_future_due_at_does_not_run_early():
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 12, 0, tzinfo=SHANGHAI_TZ).timestamp()
            due = now + 600
            mark_failed(
                repo, "snapshot", due_at=due, now=now,
                stat_date="2026-11-03", error="boom",
            )
            nxt = startup_next_run(
                repo, "snapshot",
                wall_due_at=next_daily_run(now, "23:15"),
                now=now,
            )
            assert nxt == due
            clock = [now]
            calls = []
            sched = Scheduler(
                [
                    ScheduledTask(
                        "snapshot", 86400.0, lambda: calls.append(1), next_run=nxt,
                    )
                ],
                clock=lambda: clock[0],
            )
            assert sched.run_due() == []
            assert calls == []
        finally:
            db.close()


def test_due_at_passed_runs_including_failed_retry():
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 12, 10, tzinfo=SHANGHAI_TZ).timestamp()
            due = now - 30  # 已过重试点
            mark_failed(
                repo, "snapshot", due_at=due, now=due - 300,
                stat_date="2026-11-03", error="boom",
            )
            nxt = startup_next_run(
                repo, "snapshot",
                wall_due_at=next_daily_run(now, "23:15"),
                now=now,
            )
            # 不得退回今晚墙钟；应立刻可跑
            assert nxt == due
            assert nxt <= now
            wall = next_daily_run(now, "23:15")
            assert nxt != wall
            calls = []
            clock = [now]
            sched = Scheduler(
                [
                    ScheduledTask(
                        "snapshot", 86400.0, lambda: calls.append(1), next_run=nxt,
                    )
                ],
                clock=lambda: clock[0],
            )
            results = sched.run_due()
            assert len(results) == 1
            assert calls == [1]
        finally:
            db.close()


def test_done_writes_next_due_and_blocks_until_then():
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 23, 40, tzinfo=SHANGHAI_TZ).timestamp()
            next_due = next_daily_run(now, "23:15")
            mark_done(repo, "snapshot", due_at=next_due, now=now, stat_date="2026-11-03")
            state = load_state(repo, "snapshot")
            assert state is not None
            assert state.status == STATUS_DONE
            nxt = startup_next_run(
                repo, "snapshot", wall_due_at=next_due, now=now + 60,
            )
            assert nxt == next_due
            assert nxt > now + 60
        finally:
            db.close()


def test_no_state_falls_back_to_wall_clock():
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 10, 0, tzinfo=SHANGHAI_TZ).timestamp()
            wall = next_daily_run(now, "23:15")
            nxt = startup_next_run(repo, "snapshot", wall_due_at=wall, now=now)
            assert nxt == wall
        finally:
            db.close()


def main_test():
    test_restart_mid_run_resumes_same_stat_date()
    test_seeded_before_snapshot_at_does_not_force_run()
    test_failed_future_due_at_does_not_run_early()
    test_due_at_passed_runs_including_failed_retry()
    test_done_writes_next_due_and_blocks_until_then()
    test_no_state_falls_back_to_wall_clock()
    print("OK: task_state persistence / resume / due_at gates")


if __name__ == "__main__":
    main_test()