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_DEFERRED,
    STATUS_DONE,
    STATUS_FAILED,
    STATUS_PENDING,
    STATUS_RUNNING,
    bootstrap_from_last,
    ensure_pending,
    load_state,
    mark_deferred,
    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（刻意 ≠ 墙钟）且未到不跑。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 10, 30, tzinfo=SHANGHAI_TZ).timestamp()
            # 刻意偏离墙钟 23:15，证明认落盘而非 fallback
            slot = datetime(2026, 11, 3, 22, 0, tzinfo=SHANGHAI_TZ).timestamp()
            wall = next_daily_run(now, "23:15")
            assert slot != wall
            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=wall,
                now=now,
            )
            assert nxt == slot
            assert nxt != wall
            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 test_fail_or_defer_preserves_progress():
    """H7: mark_failed / mark_deferred 不传 progress 时保留已有 progress 与 stat_date。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 12, 0, tzinfo=SHANGHAI_TZ).timestamp()
            progress = {"done": 42, "total": 100, "cursor": "aid:99"}
            mark_running(
                repo, "snapshot", stat_date="2026-11-03", now=now,
                progress=progress,
            )
            before = load_state(repo, "snapshot")
            assert before is not None
            assert before.progress == progress
            assert before.stat_date == "2026-11-03"

            due_fail = now + 300
            mark_failed(
                repo, "snapshot", due_at=due_fail, now=now + 1,
                progress=None, error="boom",
            )
            failed = load_state(repo, "snapshot")
            assert failed is not None
            assert failed.status == STATUS_FAILED
            assert failed.due_at == due_fail
            assert failed.progress == progress
            assert failed.stat_date == "2026-11-03"
            assert failed.error == "boom"

            due_defer = now + 600
            mark_deferred(
                repo, "snapshot", due_at=due_defer, now=now + 2,
                progress=None, error="snapshot_quiet_window",
            )
            deferred = load_state(repo, "snapshot")
            assert deferred is not None
            assert deferred.status == STATUS_DEFERRED
            assert deferred.due_at == due_defer
            assert deferred.progress == progress
            assert deferred.stat_date == "2026-11-03"
            assert deferred.error == "snapshot_quiet_window"
        finally:
            db.close()



def test_bootstrap_from_last_fail_due_soon():
    """A: 仅有 task:last 失败、无 state → 回填 failed 且 due_at 可立刻补跑。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 12, 0, tzinfo=SHANGHAI_TZ).timestamp()
            wall = next_daily_run(now, "23:15")
            repo.kv_set(
                "task:last:snapshot",
                json.dumps({
                    "name": "snapshot",
                    "ok": False,
                    "skipped": False,
                    "started_at": int(now - 600),
                    "finished_at": int(now - 500),
                    "error": "boom",
                    "summary": {"stat_date": "2026-11-03"},
                }, ensure_ascii=False),
            )
            assert load_state(repo, "snapshot") is None
            boot = bootstrap_from_last(
                repo, "snapshot", wall_due_at=wall, now=now, retry_seconds=300.0,
            )
            assert boot is not None
            assert boot.status == STATUS_FAILED
            assert boot.stat_date == "2026-11-03"
            assert boot.due_at <= now
            nxt = startup_next_run(
                repo, "snapshot", wall_due_at=wall, now=now, slot_hhmm="23:15",
            )
            assert nxt <= now
            assert nxt != wall
        finally:
            db.close()


def test_ensure_pending_does_not_glue_yesterday_stat_date():
    """B: 昨天未完成保留；不得把今晚 due 粘到昨天的 stat_date 上。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 0, 10, tzinfo=SHANGHAI_TZ).timestamp()
            yday_due = now - 60
            ensure_pending(
                repo, "snapshot", due_at=yday_due,
                stat_date="2026-11-02", now=now - 120,
            )
            tonight = datetime(2026, 11, 3, 23, 15, tzinfo=SHANGHAI_TZ).timestamp()
            # seed 想挂今晚：不得冲掉昨天未完成
            kept = ensure_pending(
                repo, "snapshot", due_at=tonight,
                stat_date="2026-11-03", now=now,
            )
            assert kept.stat_date == "2026-11-02"
            assert kept.due_at == yday_due
            assert resume_stat_date(
                repo, "snapshot", fallback="2026-11-03",
            ) == "2026-11-02"
        finally:
            db.close()


def test_cross_midnight_mid_seal_resumes_yesterday():
    """规则1 / B: 23:15 崩溃 → 00:10 重启续跑昨天，due 立刻可跑。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            crash = datetime(2026, 11, 2, 23, 40, tzinfo=SHANGHAI_TZ).timestamp()
            mark_running(
                repo, "snapshot", stat_date="2026-11-02", now=crash,
                progress={"done": 3, "total": 10},
            )
            repo.kv_set("snapshot:seeded:2026-11-02", "1")
            restart = datetime(2026, 11, 3, 0, 10, tzinfo=SHANGHAI_TZ).timestamp()
            wall = next_daily_run(restart, "23:15")
            nxt = startup_next_run(
                repo, "snapshot", wall_due_at=wall, now=restart, slot_hhmm="23:15",
            )
            assert nxt <= restart
            assert resume_stat_date(
                repo, "snapshot", fallback="2026-11-03",
            ) == "2026-11-02"
            state = load_state(repo, "snapshot")
            assert state is not None
            assert state.status == STATUS_PENDING
            assert state.progress == {"done": 3, "total": 10}
        finally:
            db.close()


def test_scheduler_due_reader_picks_up_enqueued_seal():
    """D: 同进程 seed 写盘 due_at=now 后，下一轮 run_due 须跑 seal。"""
    with tempfile.TemporaryDirectory() as tmp:
        db, repo = _repo(tmp)
        try:
            now = datetime(2026, 11, 3, 23, 20, tzinfo=SHANGHAI_TZ).timestamp()
            # snapshot 原 next_run 挂次日；盘上被 enqueue 成 now
            tomorrow = next_daily_run(now, "23:15")
            calls = []

            def run_snap():
                calls.append("snap")
                return None

            def due_reader(name: str):
                st = load_state(repo, name)
                return float(st.due_at) if st and st.due_at else None

            clock = [now]
            sched = Scheduler(
                [ScheduledTask("snapshot", 86400.0, run_snap, next_run=tomorrow)],
                clock=lambda: clock[0],
                due_reader=due_reader,
            )
            assert sched.run_due() == []
            assert calls == []
            ensure_pending(
                repo, "snapshot", due_at=now, stat_date="2026-11-03", now=now,
            )
            results = sched.run_due()
            assert len(results) == 1
            assert calls == ["snap"]
        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()
    test_fail_or_defer_preserves_progress()
    test_bootstrap_from_last_fail_due_soon()
    test_ensure_pending_does_not_glue_yesterday_stat_date()
    test_cross_midnight_mid_seal_resumes_yesterday()
    test_scheduler_due_reader_picks_up_enqueued_seal()
    print("OK: task_state persistence / resume / due_at gates")


if __name__ == "__main__":
    main_test()