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()