Initial import of Shirobana bot backend
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
This commit is contained in:
RRRosmontis
@@ -0,0 +1,314 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user