shirobana-bot/tests/test_notifications.py

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