import asyncio from concurrent.futures import ThreadPoolExecutor import os import sqlite3 import tempfile import threading import unittest from pathlib import Path import sys from unittest.mock import patch BASE = Path(__file__).resolve().parents[1] sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)] import channel import identity # Isolate the module-level NotificationCenter created during import as well as # the per-test centers; restore the environment immediately for other modules. _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") import notifications if _OLD_DB_PATH is None: os.environ.pop("AGENT_DB_PATH", None) else: os.environ["AGENT_DB_PATH"] = _OLD_DB_PATH from agent.modes import ModeStore def tearDownModule(): _IMPORT_TMP.cleanup() class FakeChannel(channel.Channel): def __init__(self, name, ok=True): self.name = name self.ok = ok self.sent = [] async def send(self, msg_type, target_id, content): if self.ok: self.sent.append((msg_type, target_id, content)) return self.ok class NotificationCenterTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.tmp = tempfile.TemporaryDirectory() identity.configure(os.path.join(self.tmp.name, "accounts"), admin_tg_ids=[]) self.uid = 9000000001 self.identity_patcher = patch.object(identity, "push_targets", return_value=[("tg", 10000000001)]) self.identity_patcher.start() self.center = notifications.NotificationCenter(os.path.join(self.tmp.name, "events.db")) self.modes = ModeStore(os.path.join(self.tmp.name, "events.db")) self.tg = FakeChannel("tg") channel.register_channel(self.tg) self.acks = [] self.center.register_acknowledger("mail", self.acks.append) async def asyncTearDown(self): channel.unregister_channel("tg") self.identity_patcher.stop() self.tmp.cleanup() def event(self): return notifications.Event( kind="mail", dedupe_key="mail:1", uid=self.uid, legacy_text="#2645 raw", payload={"mid": 2645, "addr": "a@example.test", "url": "https://mail.example.test"}, ) async def test_chat_render_delivery_ack_and_dedupe(self): self.modes.set_mode("tg", self.uid, "chat") async def renderer(event, mode, target): return "主人，邮件 2645 到了：https://mail.example.test" self.center.configure(renderer) with patch("agent.modes.ModeStore", return_value=self.modes): event_id = self.center.publish(self.event()) self.assertEqual(self.center.publish(self.event()), event_id) count = await self.center.deliver_due() self.assertEqual(count, 1) self.assertEqual(len(self.tg.sent), 1) self.assertIn("主人", self.tg.sent[0][2]) self.assertEqual(len(self.acks), 1) async def test_command_mode_uses_legacy_and_retries(self): self.tg.ok = False with patch("agent.modes.ModeStore", return_value=self.modes): self.center.publish(self.event()) self.assertEqual(await self.center.deliver_due(), 0) self.tg.ok = True # Make the persisted row immediately due again. import sqlite3 with sqlite3.connect(self.center.db_path) as con: con.execute("UPDATE notification_deliveries SET next_attempt_at=0") self.assertEqual(await self.center.deliver_due(), 1) self.assertEqual(self.tg.sent[0][2], "#2645 raw") async def test_chat_event_waits_for_active_conversation(self): self.modes.set_mode("tg", self.uid, "chat") async def renderer(event, mode, target): return "主人，邮件 2645 到了：https://mail.example.test" self.center.configure(renderer) with patch("agent.modes.ModeStore", return_value=self.modes): self.center.begin_conversation(self.uid) self.center.publish(self.event()) self.assertEqual(await self.center.deliver_due(), 0) self.assertEqual(self.tg.sent, []) self.center.end_conversation(self.uid) self.assertEqual(await self.center.deliver_due(), 1) self.assertEqual(len(self.tg.sent), 1) async def test_zero_target_publish_is_rejected_and_repeat_adds_bindings(self): with patch.object(identity, "push_targets", return_value=[]): self.assertIsNone(self.center.publish(self.event())) with sqlite3.connect(self.center.db_path) as con: self.assertEqual(con.execute("SELECT COUNT(*) FROM notification_events").fetchone()[0], 1) self.assertEqual(con.execute("SELECT COUNT(*) FROM notification_deliveries").fetchone()[0], 0) event_id = self.center.publish(self.event()) self.assertIsNotNone(event_id) with patch.object( identity, "push_targets", return_value=[("tg", 10000000001), ("qq", 20002)], ): self.assertEqual(self.center.publish(self.event()), event_id) with sqlite3.connect(self.center.db_path) as con: targets = con.execute( "SELECT channel,target_id FROM notification_deliveries ORDER BY channel" ).fetchall() self.assertEqual(targets, [("qq", 20002), ("tg", 10000000001)]) async def test_same_key_concurrent_publish_uses_one_event(self): event = notifications.Event( kind="mail", dedupe_key="mail:concurrent", legacy_text="raw", channel="tg", msg_type="private", target_id=10001, ) other = notifications.NotificationCenter(self.center.db_path) barrier = threading.Barrier(2) def publish(center): barrier.wait() return center.publish(event) with patch.object(notifications.NotificationCenter, "_mode", return_value="command"): with ThreadPoolExecutor(max_workers=2) as pool: ids = list(pool.map(publish, (self.center, other))) self.assertEqual(ids[0], ids[1]) with sqlite3.connect(self.center.db_path) as con: self.assertEqual(con.execute("SELECT COUNT(*) FROM notification_events").fetchone()[0], 1) self.assertEqual(con.execute("SELECT COUNT(*) FROM notification_deliveries").fetchone()[0], 1) async def test_concurrent_post_chat_flush_claims_delivery_once(self): self.center.publish(self.event()) other = notifications.NotificationCenter(self.center.db_path) entered = asyncio.Event() release = asyncio.Event() async def controlled_send(ch, msg_type, target, content, mirror=False): entered.set() await release.wait() self.tg.sent.append((msg_type, target, content)) return True with patch("notifications.send_message", side_effect=controlled_send): first = asyncio.create_task(self.center.deliver_due()) await asyncio.wait_for(entered.wait(), timeout=2) self.assertEqual(await other.deliver_due(), 0) release.set() self.assertEqual(await first, 1) self.assertEqual(len(self.tg.sent), 1) self.assertEqual(len(self.acks), 1) async def test_expired_claim_lease_recovers_after_worker_crash(self): self.center.publish(self.event()) with sqlite3.connect(self.center.db_path) as con: con.execute( "UPDATE notification_deliveries SET claim_owner='crashed',claim_until=0" ) self.assertEqual(await self.center.deliver_due(), 1) self.assertEqual(len(self.tg.sent), 1) with sqlite3.connect(self.center.db_path) as con: owner, until = con.execute( "SELECT claim_owner,claim_until FROM notification_deliveries" ).fetchone() self.assertIsNone(owner) self.assertIsNone(until) async def test_disabled_target_is_suppressed_and_republish_reenables(self): event_id = self.center.publish(self.event()) with patch.object(identity, "push_targets", return_value=[]): self.assertEqual(await self.center.deliver_due(), 0) self.assertEqual(self.tg.sent, []) with sqlite3.connect(self.center.db_path) as con: status = con.execute("SELECT status FROM notification_deliveries").fetchone()[0] self.assertEqual(status, "suppressed") self.assertEqual(self.center.publish(self.event()), event_id) self.assertEqual(await self.center.deliver_due(), 1) self.assertEqual(len(self.tg.sent), 1) async def test_active_head_does_not_block_later_deliverable_event(self): first = notifications.Event( kind="mail", dedupe_key="mail:first", uid=self.uid, legacy_text="first", created_at=1, ) second = notifications.Event( kind="mail", dedupe_key="mail:second", uid=self.uid + 1, legacy_text="second", created_at=2, ) self.modes.set_mode("tg", self.uid, "chat") with patch("agent.modes.ModeStore", return_value=self.modes): self.center.publish(first) self.center.publish(second) self.center.begin_conversation(self.uid) self.assertEqual(await self.center.deliver_due(limit=1), 1) self.assertEqual([item[2] for item in self.tg.sent], ["second"]) if __name__ == "__main__": unittest.main()