shirobana-bot/tests/test_scheduler.py

import os
import sqlite3
import sys
import tempfile
import unittest
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
from zoneinfo import ZoneInfo

# The module-level Scheduler must use a temporary DB. Prefer the real
# notifications module in deployment; only standalone local runs need a stub.
_IMPORT_TMP = tempfile.TemporaryDirectory()
_OLD_DB_PATH = os.environ.get("AGENT_DB_PATH")
os.environ["AGENT_DB_PATH"] = str(Path(_IMPORT_TMP.name) / "module-import.db")
BACKEND = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND / "bot-utility"))

try:
    import notifications  # noqa: F401
except ModuleNotFoundError:
    notifications_stub = ModuleType("notifications")

    @dataclass
    class FakeEvent:
        kind: str
        dedupe_key: str
        legacy_text: str
        uid: int | None = None
        payload: dict = field(default_factory=dict)
        channel: str | None = None
        msg_type: str | None = None
        target_id: int | None = None

    notifications_stub.Event = FakeEvent
    notifications_stub.center = None
    sys.modules["notifications"] = notifications_stub

from scheduler import HELP_TEXT, Scheduler, format_task

SHANGHAI = ZoneInfo("Asia/Shanghai")
NOW = datetime(2026, 10, 30, 13, 56, 16, tzinfo=SHANGHAI)


def tearDownModule():
    if _OLD_DB_PATH is None:
        os.environ.pop("AGENT_DB_PATH", None)
    else:
        os.environ["AGENT_DB_PATH"] = _OLD_DB_PATH
    _IMPORT_TMP.cleanup()


class FakeCenter:
    """Small persistent-dedupe fake; never sends a real notification."""

    def __init__(self):
        self.events = []
        self.keys = set()
        self.accept = True

    def publish(self, event):
        if not self.accept:
            return None
        if event.dedupe_key in self.keys:
            return "duplicate"
        self.keys.add(event.dedupe_key)
        self.events.append(event)
        return "new"


class SchedulerTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = str(Path(self.tmp.name) / "tasks.db")
        self.scheduler = Scheduler(self.db_path)
        self.center = FakeCenter()
        self.private = SimpleNamespace(
            user_id=10001,
            channel=SimpleNamespace(name="tg"),
            msg_type="private",
            target_id=10001,
        )
        self.other_private = SimpleNamespace(
            user_id=10002,
            channel=SimpleNamespace(name="tg"),
            msg_type="private",
            target_id=10002,
        )
        self.group = SimpleNamespace(
            user_id=10001,
            channel=SimpleNamespace(name="tg"),
            msg_type="group",
            target_id=20001,
        )

    def tearDown(self):
        self.tmp.cleanup()

    def create(self, ctx=None, **params):
        return self.scheduler.create_from_params(ctx or self.private, params, now=NOW)

    def test_once_naive_and_aware_run_at(self):
        naive = self.create(message="naive", kind="once", run_at="2026-10-30T14:00:00")
        expected_naive = datetime(2026, 10, 30, 14, 0, tzinfo=SHANGHAI).timestamp()
        self.assertEqual(naive["next_run_at"], expected_naive)
        self.assertEqual(naive["next_run_text"], "2026-10-30 14:00:00")
        self.assertEqual(len(naive["id"]), 8)
        int(naive["id"], 16)

        aware = self.create(message="aware", kind="once", run_at="2026-10-30T06:30:00+00:00")
        self.assertEqual(aware["next_run_text"], "2026-10-30 14:30:00")

    def test_daily_calculation_and_rollover(self):
        later = self.create(message="later", kind="daily", time_of_day="14:00:01")
        self.assertEqual(later["next_run_text"], "2026-10-30 14:00:01")
        self.assertEqual(later["time_of_day"], "14:00:01")

        tomorrow = self.create(message="tomorrow", kind="daily", time_of_day="13:00")
        self.assertEqual(tomorrow["next_run_text"], "2026-10-31 13:00:00")
        self.assertEqual(tomorrow["time_of_day"], "13:00:00")

    def test_weekly_calculation_and_weekday_deduplication(self):
        # 2026-10-30 is Friday (ISO weekday 5); Friday noon has passed.
        monday = self.create(
            message="weekly", kind="weekly", time_of_day="12:00", weekdays=[5, 1, 5]
        )
        self.assertEqual(monday["weekdays"], [1, 5])
        self.assertEqual(monday["next_run_text"], "2026-11-02 12:00:00")

        friday = self.create(
            message="friday", kind="weekly", time_of_day="14:00", weekdays=[5]
        )
        self.assertEqual(friday["next_run_text"], "2026-10-30 14:00:00")

    def test_invalid_params(self):
        cases = [
            {"message": "x", "kind": "once", "run_at": "2026-10-31", "extra": 1},
            {"message": "  ", "kind": "daily", "time_of_day": "12:00"},
            {"message": "x" * 1001, "kind": "daily", "time_of_day": "12:00"},
            {"message": "x", "kind": "monthly"},
            {"message": "x", "kind": "once"},
            {"message": "x", "kind": "once", "run_at": "not-a-time"},
            {"message": "x", "kind": "once", "run_at": "2026-10-30T13:56:16+08:00"},
            {"message": "x", "kind": "daily"},
            {"message": "x", "kind": "daily", "time_of_day": "24:00"},
            {"message": "x", "kind": "daily", "time_of_day": "12:60"},
            {"message": "x", "kind": "weekly", "time_of_day": "12:00"},
            {"message": "x", "kind": "weekly", "time_of_day": "12:00", "weekdays": []},
            {"message": "x", "kind": "weekly", "time_of_day": "12:00", "weekdays": [0]},
            {"message": "x", "kind": "weekly", "time_of_day": "12:00", "weekdays": [True]},
            {"message": "x", "kind": "weekly", "time_of_day": "12:00", "weekdays": "1"},
        ]
        for params in cases:
            with self.subTest(params=params):
                with self.assertRaises(ValueError):
                    self.create(**params)
        self.assertEqual(self.scheduler.list_for_context(self.private), [])

    def test_list_is_context_scoped_and_sorted(self):
        late = self.create(message="late", kind="daily", time_of_day="18:00")
        early = self.create(message="early", kind="daily", time_of_day="09:00")
        self.create(ctx=self.other_private, message="other user", kind="daily", time_of_day="10:00")
        group_task = self.create(ctx=self.group, message="group", kind="daily", time_of_day="11:00")

        tasks = self.scheduler.list_for_context(self.private)
        # Today's 18:00 is due before tomorrow's 09:00.
        self.assertEqual([task["id"] for task in tasks], [late["id"], early["id"]])
        self.assertEqual(self.scheduler.list_for_context(self.other_private)[0]["message"], "other user")
        self.assertEqual([task["id"] for task in self.scheduler.list_for_context(self.group)], [group_task["id"]])
        other_group = SimpleNamespace(
            user_id=10001, channel=SimpleNamespace(name="tg"), msg_type="group", target_id=20002
        )
        self.assertEqual(self.scheduler.list_for_context(other_group), [])

    def test_cancel_requires_unique_prefix_and_same_context(self):
        with patch("scheduler.secrets.token_hex", side_effect=["aaaa0001", "aaaa0002"]):
            first = self.create(message="first", kind="daily", time_of_day="10:00")
            second = self.create(message="second", kind="daily", time_of_day="11:00")
        with self.assertRaises(ValueError):
            self.scheduler.cancel_for_context(self.other_private, first["id"])
        with self.assertRaises(ValueError):
            self.scheduler.cancel_for_context(self.private, "aaaa")
        cancelled = self.scheduler.cancel_for_context(self.private, "aaaa0001")
        self.assertEqual(cancelled["status"], "cancelled")
        remaining = self.scheduler.list_for_context(self.private)
        self.assertEqual([task["id"] for task in remaining], [second["id"]])
        with self.assertRaises(ValueError):
            self.scheduler.cancel_for_context(self.private, first["id"])

    def test_fire_once_publishes_event_and_completes(self):
        task = self.create(message="喝水", kind="once", run_at="2026-10-30T14:00:00")
        fire_time = datetime(2026, 10, 30, 14, 0, 1, tzinfo=SHANGHAI)
        self.assertEqual(self.scheduler.fire_due(now=fire_time, center=self.center), 1)
        self.assertEqual(len(self.center.events), 1)
        event = self.center.events[0]
        self.assertEqual(event.kind, "reminder")
        self.assertEqual(event.uid, self.private.user_id)
        self.assertEqual(event.legacy_text, "定时提醒：喝水")
        self.assertEqual(event.payload["task_id"], task["id"])
        self.assertEqual(event.payload["message"], "喝水")
        self.assertEqual(event.payload["kind"], "once")
        self.assertEqual(event.payload["scheduled_at"], task["next_run_at"])
        self.assertEqual(event.dedupe_key, f"reminder:{task['id']}:{int(task['next_run_at'])}")
        self.assertEqual(self.scheduler.list_for_context(self.private), [])
        with sqlite3.connect(self.db_path) as con:
            status = con.execute("SELECT status FROM scheduled_tasks WHERE id=?", (task["id"],)).fetchone()[0]
        self.assertEqual(status, "completed")
        self.assertEqual(self.scheduler.fire_due(now=fire_time, center=self.center), 0)
        self.assertEqual(len(self.center.events), 1)

    def test_fire_group_event_uses_group_target(self):
        self.create(ctx=self.group, message="群提醒", kind="once", run_at="2026-10-30T14:00:00")
        self.scheduler.fire_due(now=NOW.timestamp() + 600, center=self.center)
        event = self.center.events[0]
        self.assertIsNone(event.uid)
        self.assertEqual((event.channel, event.msg_type, event.target_id), ("tg", "group", 20001))

    def test_repeating_tasks_advance_past_now_without_catchup_burst(self):
        daily = self.create(message="daily", kind="daily", time_of_day="14:00")
        weekly = self.create(
            message="weekly", kind="weekly", time_of_day="14:00", weekdays=[5]
        )
        # Several daily occurrences and the first weekly occurrence were missed.
        later = datetime(2026, 11, 2, 15, 0, tzinfo=SHANGHAI)
        self.assertEqual(self.scheduler.fire_due(now=later, center=self.center), 2)
        self.assertEqual(len(self.center.events), 2)
        tasks = {task["id"]: task for task in self.scheduler.list_for_context(self.private)}
        self.assertEqual(tasks[daily["id"]]["next_run_text"], "2026-11-03 14:00:00")
        self.assertEqual(tasks[weekly["id"]]["next_run_text"], "2026-11-06 14:00:00")
        self.assertEqual(self.scheduler.fire_due(now=later, center=self.center), 0)
        self.assertEqual(len(self.center.events), 2)

    def test_restart_persistence_and_publish_dedupe_after_crash_window(self):
        task = self.create(message="persist", kind="once", run_at="2026-10-30T14:00:00")
        due = datetime(2026, 10, 30, 14, 0, 1, tzinfo=SHANGHAI)
        restarted = Scheduler(self.db_path)
        self.assertEqual([item["id"] for item in restarted.list_for_context(self.private)], [task["id"]])
        restarted.fire_due(now=due, center=self.center)
        self.assertEqual(len(self.center.events), 1)

        # Simulate a crash after publish but before the task status commit.
        with sqlite3.connect(self.db_path) as con:
            con.execute(
                "UPDATE scheduled_tasks SET status='active',next_run_at=? WHERE id=?",
                (task["next_run_at"], task["id"]),
            )
        restarted_again = Scheduler(self.db_path)
        restarted_again.fire_due(now=due, center=self.center)
        self.assertEqual(len(self.center.events), 1, "same dedupe key must not publish twice")
        self.assertEqual(restarted_again.list_for_context(self.private), [])

    def test_zero_delivery_does_not_advance_and_later_retry_is_accepted(self):
        task = self.create(message="等待绑定", kind="once", run_at="2026-10-30T14:00:00")
        due = datetime(2026, 10, 30, 14, 0, 1, tzinfo=SHANGHAI)
        self.center.accept = False
        self.assertEqual(self.scheduler.fire_due(now=due, center=self.center), 0)
        self.assertEqual(self.center.events, [])
        remaining = self.scheduler.list_for_context(self.private)
        self.assertEqual([item["id"] for item in remaining], [task["id"]])
        self.assertEqual(remaining[0]["next_run_at"], task["next_run_at"])

        self.center.accept = True
        self.assertEqual(self.scheduler.fire_due(now=due, center=self.center), 1)
        self.assertEqual(len(self.center.events), 1)
        self.assertEqual(self.scheduler.list_for_context(self.private), [])

    def test_cancel_after_select_before_claim_prevents_publish(self):
        task = self.create(message="不要发送", kind="once", run_at="2026-10-30T14:00:00")
        due = datetime(2026, 10, 30, 14, 0, 1, tzinfo=SHANGHAI)
        original_claim = self.scheduler._claim_task

        def cancel_then_claim(row, current, owner):
            self.scheduler.cancel_for_context(self.private, row["id"])
            return original_claim(row, current, owner)

        with patch.object(self.scheduler, "_claim_task", side_effect=cancel_then_claim):
            self.assertEqual(self.scheduler.fire_due(now=due, center=self.center), 0)
        self.assertEqual(self.center.events, [])
        with sqlite3.connect(self.db_path) as con:
            status = con.execute(
                "SELECT status FROM scheduled_tasks WHERE id=?", (task["id"],)
            ).fetchone()[0]
        self.assertEqual(status, "cancelled")

    def test_expired_firing_lease_recovers_without_duplicate_event(self):
        task = self.create(message="崩溃恢复", kind="once", run_at="2026-10-30T14:00:00")
        due = datetime(2026, 10, 30, 14, 0, 1, tzinfo=SHANGHAI)
        with sqlite3.connect(self.db_path) as con:
            con.execute(
                """UPDATE scheduled_tasks
                   SET status='firing',firing_owner='crashed',firing_until=0
                   WHERE id=?""",
                (task["id"],),
            )
        self.assertEqual(self.scheduler.fire_due(now=due, center=self.center), 1)
        self.assertEqual(len(self.center.events), 1)
        self.assertEqual(self.scheduler.list_for_context(self.private), [])

    def test_format_task_and_help_text(self):
        task = self.create(message="写周报", kind="weekly", time_of_day="18:00", weekdays=[1, 5])
        text = format_task(task)
        self.assertIn(task["id"], text)
        self.assertIn("每周", text)
        self.assertIn("写周报", text)
        self.assertIn(task["next_run_text"], text)
        self.assertIn("自然语言", HELP_TEXT)
        self.assertIn("查看", HELP_TEXT)
        self.assertIn("取消", HELP_TEXT)


if __name__ == "__main__":
    unittest.main()