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,221 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user