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,102 @@
|
||||
import asyncio
|
||||
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 agent_loop
|
||||
from agent import ToolResult
|
||||
|
||||
|
||||
class Channel:
|
||||
name = "tg"
|
||||
def is_admin(self, uid):
|
||||
return True
|
||||
|
||||
|
||||
class Store:
|
||||
instances = []
|
||||
def __init__(self):
|
||||
self.steps = []
|
||||
self.finished = None
|
||||
Store.instances.append(self)
|
||||
def start(self, uid, conversation):
|
||||
return "task"
|
||||
def record(self, *args):
|
||||
self.steps.append(args)
|
||||
def finish(self, task, status, error=""):
|
||||
self.finished = status
|
||||
|
||||
|
||||
class Executor:
|
||||
calls = []
|
||||
def __init__(self):
|
||||
pass
|
||||
def prompt_schema(self, ctx):
|
||||
return []
|
||||
async def execute(self, name, args, ctx):
|
||||
Executor.calls.append((name, args))
|
||||
if name == "setu":
|
||||
return ToolResult(True, "ok", "图片已获取", artifacts=[{"id": "a1"}])
|
||||
return ToolResult(True, "ok", "邮件已发送", data={"used": args})
|
||||
|
||||
|
||||
class Engine:
|
||||
def __init__(self, responses):
|
||||
self.responses = list(responses)
|
||||
async def call_deepseek(self, *args, **kwargs):
|
||||
return self.responses.pop(0)
|
||||
def parse_json_response(self, raw):
|
||||
return raw
|
||||
async def run_tools(self, calls):
|
||||
return "web result"
|
||||
|
||||
|
||||
class AgentLoopTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_two_round_composition_and_final_response(self):
|
||||
Executor.calls = []
|
||||
initial = {"message": [], "tools": [{"name": "setu", "args": {"action": "get", "params": {"tags": ["铃兰"]}}}]}
|
||||
second = {"message": [], "tools": [{"name": "mail", "args": {"action": "send", "params": {"artifact_ids": ["a1"]}}}]}
|
||||
final = {"message": [{"content": "完成了"}], "tools": []}
|
||||
engine = Engine([second, final])
|
||||
with patch.object(agent_loop, "AgentTaskStore", Store), \
|
||||
patch.object(agent_loop, "ToolExecutor", Executor):
|
||||
result = await agent_loop.run_agent_loop(
|
||||
engine, initial, "system", "user", Channel(), "private", 1, 2)
|
||||
self.assertEqual(result, final)
|
||||
self.assertEqual([x[0] for x in Executor.calls], ["setu", "mail"])
|
||||
self.assertEqual(Store.instances[-1].finished, "done")
|
||||
self.assertEqual(len(Store.instances[-1].steps), 2)
|
||||
|
||||
async def test_not_understood_never_executes(self):
|
||||
initial = {"understood": False, "message": [], "tools": [{"name": "mail"}]}
|
||||
with patch.object(agent_loop, "ToolExecutor", Executor):
|
||||
result = await agent_loop.run_agent_loop(
|
||||
Engine([]), initial, "s", "u", Channel(), "private", 1, 2)
|
||||
self.assertEqual(result["tools"], [])
|
||||
self.assertIn("明白", result["message"][0]["content"])
|
||||
self.assertEqual(Executor.calls, [])
|
||||
|
||||
async def test_model_failure_includes_specific_safe_error(self):
|
||||
Executor.calls = []
|
||||
initial = {"message": [], "tools": [{"name": "pixiv", "args": {"action": "get"}}]}
|
||||
with patch.object(agent_loop, "AgentTaskStore", Store), \
|
||||
patch.object(agent_loop, "ToolExecutor", Executor):
|
||||
# Force an observation containing a raw-looking failure, then no model response.
|
||||
original = Executor.execute
|
||||
async def failure(self, name, args, ctx):
|
||||
return ToolResult(False, "send_failed", "多次尝试发送失败。请手动访问 https://example.test/image")
|
||||
with patch.object(Executor, "execute", failure):
|
||||
result = await agent_loop.run_agent_loop(
|
||||
Engine([None]), initial, "s", "u", Channel(), "private", 1, 2)
|
||||
text = result["message"][0]["content"]
|
||||
self.assertIn("多次尝试发送失败", text)
|
||||
self.assertIn("send_failed", text)
|
||||
self.assertIn("https://example.test/image", text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,417 @@
|
||||
"""Isolated tests for the agent tool runtime; no real API or message delivery."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from agent import ArtifactError, ArtifactStore, CaptureChannel, ModeStore, ToolContext, ToolExecutor
|
||||
import channel as channel_module
|
||||
|
||||
|
||||
class FakeChannel:
|
||||
name = "qq"
|
||||
|
||||
def __init__(self):
|
||||
self.sent = []
|
||||
|
||||
async def send(self, msg_type, target_id, content):
|
||||
self.sent.append((msg_type, target_id, content))
|
||||
return True
|
||||
|
||||
|
||||
class ModeStoreTests(unittest.TestCase):
|
||||
def test_default_private_and_group_scopes(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
store = ModeStore(Path(directory) / "agent.db")
|
||||
self.assertEqual(store.get_mode("qq", 10), "command")
|
||||
store.set_mode("qq", 10, "agent")
|
||||
store.set_mode("qq", 10, "agent", group_id=20)
|
||||
store.set_mode("tg", 10, "agent")
|
||||
self.assertEqual(store.get_mode("qq", 10), "agent")
|
||||
self.assertEqual(store.get_mode("qq", 999, group_id=20), "agent")
|
||||
self.assertEqual(store.get_mode("qq", 10, group_id=21), "command")
|
||||
store.set_mode("qq", 10, "command", group_id=20)
|
||||
self.assertEqual(store.get_mode("qq", 10, group_id=20), "command")
|
||||
with self.assertRaises(ValueError):
|
||||
store.set_mode("qq", 10, "invalid")
|
||||
|
||||
|
||||
class ArtifactStoreTests(unittest.TestCase):
|
||||
def test_local_base64_and_confined_materialize(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory) / "artifacts"
|
||||
source = Path(directory) / "source.jpg"
|
||||
source.write_bytes(b"image-data")
|
||||
store = ArtifactStore(root)
|
||||
local = store.register(source, "image")
|
||||
inline = store.register(base64.b64encode(b"second").decode(), "image", "two.jpg")
|
||||
self.assertNotEqual(local["artifact_id"], inline["artifact_id"])
|
||||
self.assertGreaterEqual(len(local["artifact_id"]), 16)
|
||||
first_path = store.materialize(local["artifact_id"])
|
||||
self.assertEqual(first_path.read_bytes(), b"image-data")
|
||||
self.assertTrue(first_path.is_relative_to(root))
|
||||
source.unlink()
|
||||
self.assertEqual(first_path.read_bytes(), b"image-data")
|
||||
self.assertNotIn(str(source), json.dumps(local))
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.materialize("../../etc/passwd")
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.register("/definitely/not/a/real/file")
|
||||
|
||||
def test_url_materializes_under_root(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
store = ArtifactStore(Path(directory) / "store")
|
||||
artifact = store.register("https://example.invalid/a.jpg", "image")
|
||||
response = mock.MagicMock()
|
||||
response.__enter__.return_value = response
|
||||
response.geturl.return_value = "https://example.invalid/a.jpg"
|
||||
response.headers = {}
|
||||
response.read.return_value = b"remote-image"
|
||||
with mock.patch("urllib.request.urlopen", return_value=response) as opener:
|
||||
path = store.materialize(artifact["artifact_id"])
|
||||
self.assertEqual(path.read_bytes(), b"remote-image")
|
||||
self.assertTrue(path.is_relative_to(store.root))
|
||||
opener.assert_called_once()
|
||||
|
||||
def test_scoped_artifacts_require_matching_owner_and_session(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
store = ArtifactStore(Path(directory) / "store")
|
||||
channel = SimpleNamespace(name="qq")
|
||||
ctx = SimpleNamespace(channel=channel, msg_type="group", user_id=10, target_id=20)
|
||||
artifact = store.register(b64("private"), "file", "private.txt", ctx=ctx)
|
||||
artifact_id = artifact["artifact_id"]
|
||||
self.assertEqual(store.get(artifact_id, ctx)["artifact_id"], artifact_id)
|
||||
self.assertEqual(store.materialize(artifact_id, ctx).read_bytes(), b"private")
|
||||
record = json.loads((store.root / artifact_id / "artifact.json").read_text())
|
||||
self.assertEqual(record["owner_id"], "10")
|
||||
self.assertEqual(record["session"], {
|
||||
"channel": "qq", "msg_type": "group", "target_id": "20",
|
||||
})
|
||||
foreign_contexts = [
|
||||
SimpleNamespace(channel=channel, msg_type="group", user_id=11, target_id=20),
|
||||
SimpleNamespace(channel=channel, msg_type="group", user_id=10, target_id=21),
|
||||
SimpleNamespace(channel=channel, msg_type="private", user_id=10, target_id=20),
|
||||
SimpleNamespace(channel=SimpleNamespace(name="tg"), msg_type="group", user_id=10, target_id=20),
|
||||
]
|
||||
for foreign in foreign_contexts:
|
||||
with self.subTest(foreign=foreign):
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.get(artifact_id, foreign)
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.materialize(artifact_id, foreign)
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.get(artifact_id)
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.materialize(artifact_id)
|
||||
|
||||
def test_legacy_artifacts_are_denied_to_agent_contexts(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
store = ArtifactStore(Path(directory) / "store")
|
||||
legacy = store.register(b64("legacy"), "file")
|
||||
ctx = SimpleNamespace(
|
||||
channel=SimpleNamespace(name="qq"), msg_type="private",
|
||||
user_id=10, target_id=10,
|
||||
)
|
||||
artifact_id = legacy["artifact_id"]
|
||||
self.assertEqual(store.materialize(artifact_id).read_bytes(), b"legacy")
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.get(artifact_id, ctx)
|
||||
with self.assertRaises(ArtifactError):
|
||||
store.materialize(artifact_id, ctx)
|
||||
|
||||
|
||||
class ExecutorTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.temporary = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self.temporary.name)
|
||||
self.store = ArtifactStore(self.root / "artifacts")
|
||||
self.real_channel = FakeChannel()
|
||||
self.sender = mock.AsyncMock(return_value=True)
|
||||
self.ctx = ToolContext(
|
||||
self.real_channel, "private", 9000000001, 123,
|
||||
is_admin=False, request_id="req", task_id="task",
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
self.temporary.cleanup()
|
||||
|
||||
def executor(self, modules):
|
||||
return ToolExecutor(self.store, modules=modules, send_message=self.sender)
|
||||
|
||||
async def test_capture_channel_never_delivers_or_mirrors_private_text(self):
|
||||
capture = CaptureChannel(self.ctx, self.store)
|
||||
# Use the production sender to prove its private mirror path is skipped.
|
||||
with mock.patch.object(channel_module, "_mirror_private", new=mock.AsyncMock()) as mirror:
|
||||
result = await channel_module.send_message(capture, "private", 123, "captured")
|
||||
self.assertFalse(result)
|
||||
self.assertEqual(capture.texts, ["captured"])
|
||||
self.assertEqual(self.real_channel.sent, [])
|
||||
mirror.assert_not_awaited()
|
||||
|
||||
async def test_setu_media_is_delivered_once_and_text_is_model_facing(self):
|
||||
image = self.root / "captured.jpg"
|
||||
image.write_bytes(b"jpeg")
|
||||
calls = []
|
||||
|
||||
async def handle_setu(tags, r18, ai, msg_type, target_id, ws):
|
||||
from channel import send_message
|
||||
calls.append((tags, r18, ai))
|
||||
await send_message(ws, msg_type, target_id, "正在获取图片")
|
||||
await send_message(ws, msg_type, target_id, [
|
||||
{"type": "text", "data": {"text": "PID 42"}},
|
||||
{"type": "image", "data": {"file": str(image)}},
|
||||
])
|
||||
|
||||
module = SimpleNamespace(handle_setu=handle_setu)
|
||||
executor = self.executor({"setu": module})
|
||||
result = await executor.execute(
|
||||
"setu", {"action": "get", "params": {"tags": ["猫耳"], "r18": False}}, self.ctx,
|
||||
)
|
||||
self.assertTrue(result.ok, result)
|
||||
self.assertEqual(calls, [(["猫耳"], "不要r18", None)])
|
||||
self.assertEqual(result.safe_message, "正在获取图片\nPID 42")
|
||||
self.assertEqual(len(result.artifacts), 1)
|
||||
self.assertEqual(result.data["deliveries"][0]["ok"], True)
|
||||
self.sender.assert_awaited_once()
|
||||
sent = self.sender.await_args.args[3]
|
||||
self.assertEqual(sent[0]["type"], "image")
|
||||
delivered = Path(sent[0]["data"]["file"].removeprefix("file://"))
|
||||
self.assertTrue(delivered.is_relative_to(self.store.root))
|
||||
self.assertEqual(self.real_channel.sent, [])
|
||||
|
||||
async def test_other_tool_text_is_not_delivered_and_mail_quotes_are_preserved(self):
|
||||
seen = []
|
||||
|
||||
async def handle_mail(ws, msg_type, target_id, user_id, raw):
|
||||
from channel import send_message
|
||||
seen.append(raw)
|
||||
await send_message(ws, msg_type, target_id, "邮件已发送")
|
||||
|
||||
executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
|
||||
result = await executor.execute("mail", {
|
||||
"action": "new",
|
||||
"params": {"to_address": "to@example.com", "subject": "two words", "body_text": "hello world"},
|
||||
}, self.ctx)
|
||||
self.assertTrue(result.ok)
|
||||
tokens = shlex.split(seen[0])
|
||||
self.assertEqual(tokens[-3:], ["to@example.com", "two words", "hello world"])
|
||||
self.assertEqual(result.safe_message, "邮件已发送")
|
||||
self.sender.assert_not_awaited()
|
||||
self.assertEqual(self.real_channel.sent, [])
|
||||
|
||||
async def test_mail_send_unknown_is_not_retryable(self):
|
||||
async def handle_mail(ws, msg_type, target_id, user_id, raw):
|
||||
from channel import send_message
|
||||
await send_message(ws, msg_type, target_id, "发送失败：上游返回了无法识别的响应")
|
||||
|
||||
executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
|
||||
result = await executor.execute("mail", {
|
||||
"action": "send",
|
||||
"params": {"to_address": "to@example.com", "subject": "s", "body_text": "b"},
|
||||
}, self.ctx)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.code, "send_unknown")
|
||||
self.assertFalse(result.retryable)
|
||||
|
||||
async def test_outbound_rejects_artifact_from_another_owner(self):
|
||||
foreign_ctx = ToolContext(self.real_channel, "private", self.ctx.user_id + 1, 123)
|
||||
foreign = self.store.register(b64("secret"), "file", "secret.txt", ctx=foreign_ctx)
|
||||
executor = self.executor({})
|
||||
result = await executor.execute(
|
||||
"send_file", {"artifact_id": foreign["artifact_id"]}, self.ctx,
|
||||
)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.code, "invalid_artifact")
|
||||
self.sender.assert_not_awaited()
|
||||
self.assertEqual(self.real_channel.sent, [])
|
||||
|
||||
async def test_mail_attachments_reject_foreign_artifact_before_sending(self):
|
||||
foreign_ctx = ToolContext(self.real_channel, "private", self.ctx.user_id, 999)
|
||||
foreign = self.store.register(b64("secret"), "file", "secret.txt", ctx=foreign_ctx)
|
||||
send_with_attachments = mock.AsyncMock(return_value={"ok": True})
|
||||
module = SimpleNamespace(
|
||||
_load_user=lambda uid: {
|
||||
"default": "from@rrrosmontis.icu",
|
||||
"addresses": ["from@rrrosmontis.icu"],
|
||||
},
|
||||
send_with_attachments=send_with_attachments,
|
||||
)
|
||||
executor = self.executor({"mail": module})
|
||||
result = await executor.execute("mail", {
|
||||
"action": "send",
|
||||
"params": {
|
||||
"to_address": "to@example.com", "subject": "s", "body_text": "b",
|
||||
"artifact_ids": [foreign["artifact_id"]],
|
||||
},
|
||||
}, self.ctx)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.code, "invalid_artifact")
|
||||
send_with_attachments.assert_not_awaited()
|
||||
|
||||
async def test_mail_list_limit_and_info_mid_aliases(self):
|
||||
seen = []
|
||||
async def handle_mail(ws, msg_type, target_id, user_id, raw):
|
||||
seen.append(raw)
|
||||
executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
|
||||
result = await executor.execute("mail", {"action": "list", "params": {"limit": 5}}, self.ctx)
|
||||
self.assertTrue(result.ok, result)
|
||||
result = await executor.execute("mail", {"action": "info", "params": {"mid": 2645}}, self.ctx)
|
||||
self.assertTrue(result.ok, result)
|
||||
self.assertEqual(seen, [".mail list 5", ".mail info 2645"])
|
||||
|
||||
async def test_mail_attachments_use_new_api_and_only_artifact_paths(self):
|
||||
artifact = self.store.register(b64("attachment"), "file", "a.txt", ctx=self.ctx)
|
||||
calls = []
|
||||
|
||||
async def send_with_attachments(*args):
|
||||
calls.append(args)
|
||||
return {"ok": True, "attachment_count": 1}
|
||||
|
||||
async def handle_mail(*args):
|
||||
raise AssertionError("legacy mail handler must not be called")
|
||||
|
||||
module = SimpleNamespace(
|
||||
_load_user=lambda uid: {
|
||||
"default": "from@rrrosmontis.icu",
|
||||
"addresses": ["from@rrrosmontis.icu"],
|
||||
},
|
||||
send_with_attachments=send_with_attachments,
|
||||
handle_mail=handle_mail,
|
||||
)
|
||||
executor = self.executor({"mail": module})
|
||||
result = await executor.execute("mail", {
|
||||
"action": "send",
|
||||
"params": {
|
||||
"to_address": "to@example.com", "subject": "s", "body_text": "b",
|
||||
"artifact_ids": [artifact["artifact_id"]],
|
||||
},
|
||||
}, self.ctx)
|
||||
self.assertTrue(result.ok, result)
|
||||
self.assertEqual(result.code, "sent")
|
||||
self.assertEqual(len(calls), 1)
|
||||
args = calls[0]
|
||||
self.assertEqual(args[:5], (
|
||||
self.ctx.user_id, "from@rrrosmontis.icu", "to@example.com", "s", "b",
|
||||
))
|
||||
self.assertEqual(Path(args[5][0]).read_bytes(), b"attachment")
|
||||
self.assertTrue(Path(args[5][0]).is_relative_to(self.store.root))
|
||||
self.sender.assert_not_awaited()
|
||||
|
||||
async def test_permission_is_enforced_before_admin_handler_import_or_call(self):
|
||||
handler = mock.AsyncMock()
|
||||
executor = self.executor({"server": SimpleNamespace(handle_command=handler)})
|
||||
result = await executor.execute("server", {"action": "list", "params": {}}, self.ctx)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.code, "permission_denied")
|
||||
handler.assert_not_awaited()
|
||||
schema_names = {item["name"] for item in executor.prompt_schema(self.ctx)}
|
||||
self.assertNotIn("server", schema_names)
|
||||
admin_ctx = ToolContext(self.real_channel, "group", 1, 2, is_admin=True)
|
||||
self.assertIn("server", {item["name"] for item in executor.prompt_schema(admin_ctx)})
|
||||
|
||||
async def test_send_file_accepts_only_artifact_id(self):
|
||||
executor = self.executor({})
|
||||
rejected = await executor.execute("send_file", {"artifact_id": "/etc/passwd"}, self.ctx)
|
||||
self.assertFalse(rejected.ok)
|
||||
self.assertEqual(rejected.code, "invalid_artifact")
|
||||
self.sender.assert_not_awaited()
|
||||
artifact = self.store.register(b64("file-body"), "file", "report.txt", ctx=self.ctx)
|
||||
result = await executor.execute("send_file", {"artifact_id": artifact["artifact_id"]}, self.ctx)
|
||||
self.assertTrue(result.ok, result)
|
||||
content = self.sender.await_args.args[3]
|
||||
self.assertEqual(content[0]["type"], "file")
|
||||
self.assertNotEqual(content[0]["data"]["file"], "/etc/passwd")
|
||||
self.assertTrue(Path(content[0]["data"]["file"].removeprefix("file://")).is_relative_to(self.store.root))
|
||||
|
||||
async def test_send_sticker_keeps_onebot_image_segment_for_qq(self):
|
||||
artifact = self.store.register(b64("sticker"), "image", "sticker.webp", ctx=self.ctx)
|
||||
executor = self.executor({})
|
||||
result = await executor.execute(
|
||||
"send_sticker", {"artifact_id": artifact["artifact_id"]}, self.ctx,
|
||||
)
|
||||
self.assertTrue(result.ok, result)
|
||||
content = self.sender.await_args.args[3]
|
||||
self.assertEqual(content[0]["type"], "image")
|
||||
self.assertEqual(content[0]["data"]["sub_type"], 1)
|
||||
|
||||
async def test_validation_and_handler_errors_are_safe(self):
|
||||
async def broken(*args):
|
||||
raise RuntimeError("secret=/home/private/token")
|
||||
|
||||
executor = self.executor({"log": SimpleNamespace(handle_command=broken)})
|
||||
admin = ToolContext(self.real_channel, "private", 1, 1, is_admin=True)
|
||||
result = await executor.execute("log", {"action": "tail", "params": {"lines": 5}}, admin)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertTrue(result.retryable)
|
||||
self.assertNotIn("secret", result.safe_message)
|
||||
bad = await executor.execute("log", {"action": "tail", "params": {"lines": "5;rm"}}, admin)
|
||||
self.assertEqual(bad.code, "invalid_args")
|
||||
unknown = await executor.execute("web_search", {}, admin)
|
||||
self.assertEqual(unknown.code, "unknown_tool")
|
||||
|
||||
async def test_remaining_legacy_modules_are_awaited_with_safe_commands(self):
|
||||
calls = []
|
||||
|
||||
def module(handler_name):
|
||||
async def handler(*args):
|
||||
calls.append((handler_name, args[-1]))
|
||||
return SimpleNamespace(**{handler_name: handler})
|
||||
|
||||
live = SimpleNamespace(
|
||||
HELP_TEXT="live help",
|
||||
status_text=lambda *args: "status",
|
||||
add_follow=lambda *args: calls.append(("add_follow", args)),
|
||||
remove_follow=lambda *args: None,
|
||||
)
|
||||
modules = {
|
||||
"pixiv": module("handle_pixiv"),
|
||||
"jm": module("handle_jm"),
|
||||
"ginkgo": module("handle_ginkgo"),
|
||||
"my": module("handle_my"),
|
||||
"deepseek": module("handle_command"),
|
||||
"log": module("handle_command"),
|
||||
"live": live,
|
||||
}
|
||||
executor = self.executor(modules)
|
||||
admin = ToolContext(self.real_channel, "group", 1, 2, is_admin=True)
|
||||
requests = [
|
||||
("pixiv", "get", {"pid": 123}),
|
||||
("jm", "search", {"keyword": "two words", "page": 2}),
|
||||
("ginkgo", "notice", {"count": 3}),
|
||||
("my", "push", {"platform": "qq", "enabled": True}),
|
||||
("deepseek", "alert", {"enabled": False}),
|
||||
("log", "tail", {"lines": 7}),
|
||||
("live", "follow", {"room": 4, "site": 1}),
|
||||
]
|
||||
for name, action, params in requests:
|
||||
result = await executor.execute(name, {"action": action, "params": params}, admin)
|
||||
self.assertTrue(result.ok, (name, result))
|
||||
self.assertEqual(calls[0], ("handle_pixiv", ".pixiv 123"))
|
||||
self.assertEqual(calls[1], ("handle_jm", ".jm ss two words 2"))
|
||||
self.assertEqual(calls[2], ("handle_ginkgo", ".ginkgo notice 3"))
|
||||
self.assertEqual(calls[3], ("handle_my", "qq-push on"))
|
||||
self.assertEqual(calls[4], ("handle_command", ".deepseek alert off"))
|
||||
self.assertEqual(calls[5], ("handle_command", ".log 7"))
|
||||
self.assertEqual(calls[6][0], "add_follow")
|
||||
self.assertEqual(self.real_channel.sent, [])
|
||||
self.sender.assert_not_awaited()
|
||||
|
||||
|
||||
def b64(value):
|
||||
"""Encode test payload without embedding large base64 literals."""
|
||||
return base64.b64encode(value.encode()).decode()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Akaros Mail REST and Agent tool tests; no real email is sent."""
|
||||
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from agent import ToolContext, ToolExecutor
|
||||
|
||||
mail = importlib.import_module("for-bot_mail")
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status=201, data=None):
|
||||
self.status_code = status
|
||||
self.headers = {"Content-Type": "application/json"}
|
||||
self.data = data or {"message": "accepted"}
|
||||
def json(self):
|
||||
return self.data
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
self.calls = []
|
||||
def __enter__(self):
|
||||
return self
|
||||
def __exit__(self, *args):
|
||||
return False
|
||||
def request(self, method, url, **kwargs):
|
||||
self.calls.append((method, url, kwargs))
|
||||
return self.response
|
||||
|
||||
|
||||
class AkarosMailTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_send_uses_fixed_shirobana_address(self):
|
||||
session = FakeSession(FakeResponse())
|
||||
with patch("requests.Session", return_value=session), \
|
||||
patch.object(mail, "API_KEY", "test-key"):
|
||||
result = mail.send_email(
|
||||
"user@example.test", "测试主题", "正文",
|
||||
cc=["copy@example.test"], bcc=[])
|
||||
self.assertEqual(result["message"], "accepted")
|
||||
payload = session.calls[0][2]["json"]
|
||||
self.assertEqual(payload["from"], "shirobana@akaros.si")
|
||||
self.assertEqual(payload["to"], ["user@example.test"])
|
||||
self.assertEqual(payload["cc"], ["copy@example.test"])
|
||||
self.assertEqual(payload["text"], "正文")
|
||||
|
||||
def test_invalid_recipient_is_rejected(self):
|
||||
with self.assertRaises(mail.MailInputError):
|
||||
mail.send_email("not-an-email", "subject", "text")
|
||||
|
||||
async def test_agent_send_and_result_unknown(self):
|
||||
module = SimpleNamespace(
|
||||
send_email=lambda *args: {"message": "accepted"},
|
||||
list_messages=lambda *args: [],
|
||||
get_message=lambda value: {},
|
||||
format_messages=lambda data: "没有找到邮件",
|
||||
format_message=lambda data: "正文",
|
||||
)
|
||||
executor = ToolExecutor(modules={"akaros_mail": module})
|
||||
ctx = ToolContext(channel=SimpleNamespace(name="tg"), msg_type="private", user_id=1, target_id=1)
|
||||
result = await executor.execute(
|
||||
"akaros_mail",
|
||||
{"action": "send", "params": {"to": "user@example.test", "subject": "s", "text": "t"}},
|
||||
ctx,
|
||||
)
|
||||
self.assertTrue(result.ok)
|
||||
def fail_send(*args):
|
||||
raise RuntimeError("connection lost")
|
||||
module.send_email = fail_send
|
||||
result = await executor.execute(
|
||||
"akaros_mail",
|
||||
{"action": "send", "params": {"to": "user@example.test", "subject": "s", "text": "t"}},
|
||||
ctx,
|
||||
)
|
||||
self.assertEqual(result.code, "send_unknown")
|
||||
self.assertFalse(result.retryable)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Akaros mailbox is a personal asset available without keyword activation."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from chat import engine
|
||||
|
||||
|
||||
class AkarosPersonaTests(unittest.TestCase):
|
||||
def test_mailbox_is_personal_asset_with_autonomous_tool_selection(self):
|
||||
self.assertIn("shirobana@akaros.si", engine.BASE_PROMPT)
|
||||
self.assertIn("个人资产", engine.BASE_PROMPT)
|
||||
self.assertIn("不要求用户特别指定服务名", engine.BASE_PROMPT)
|
||||
self.assertIn("不得在没有明确任务或合理上下文时", engine.BASE_PROMPT)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Expression guardrails and deterministic receipts; no live API or user data."""
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / 'bot-utility'), str(BASE)]
|
||||
import agent_loop
|
||||
from agent import ToolExecutor, ToolContext
|
||||
from chat import engine
|
||||
|
||||
|
||||
class AntiDialogueStyleTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_all_requested_counterexamples_and_templates_are_explicit(self):
|
||||
prompt = engine.load_instruction('style_instructions.txt')
|
||||
phrases = ['这句话稳', '比之前的都重', '我记下了', '我记得了', '我记住了',
|
||||
'这个账我记住了', '比之前的都轻', '睡吧，我守着', '你去，我跟',
|
||||
'问别人还是问自己', '行，xxx', '什么心思都写着xx上', '这句话接的漂亮',
|
||||
'行，我认', '这句话我不接', 'xx比xx沉', 'xx比xx重', 'xx比xx稳', 'xx比xx轻',
|
||||
'没接你的xxx', '比之前的都强', '他记得这事', '这话我留着', '很平', '稳',
|
||||
'我不拦', '我跟', '回去睡', '你这句话我记着', '认你', '我不催你',
|
||||
'这话是ta告诉你的，还是你自己想明白的', '认下来', '声音没抖，但比刚才稳',
|
||||
'比刚才轻', '比刚才重', '比刚才平', '很直', '这话我记下了', '这话我收下了',
|
||||
'说的比谁都响', '认这个', '收不回去', '认得实', 'xx就是认了', '你自己掂量',
|
||||
'没接你那句xx', '不催', '你掂量']
|
||||
for phrase in phrases:
|
||||
with self.subTest(phrase=phrase): self.assertIn(phrase, prompt)
|
||||
for principle in ('同义改写', '按共同结构识别', '即使历史里的白花曾这样说',
|
||||
'直接谈具体对象', '不是可模仿', '单个汉字全部封禁', '用户既定称呼',
|
||||
'不知道该怎么接', '不要评论自己如何接话'):
|
||||
self.assertIn(principle, prompt)
|
||||
|
||||
def test_builtin_default_is_exactly_the_deployed_expression_file(self):
|
||||
actual = engine.load_instruction('style_instructions.txt')
|
||||
self.assertEqual(actual.strip(), engine.DEFAULT_STYLE_INSTRUCTIONS.strip())
|
||||
self.assertIn('禁用：关系暗语式短应答与话语评断腔', engine.BASE_PROMPT)
|
||||
|
||||
def test_success_fallback_does_not_repeat_memory_acknowledgements(self):
|
||||
result = {'tool':'wardrobe', 'action':'add', 'ok':True,
|
||||
'data':{'item':{'id':'internal-id', 'name':'蓝色短裙', 'description':'不复述的完整细节'}}}
|
||||
output = agent_loop._message_text(agent_loop._fallback_parsed([result]))
|
||||
self.assertIn('蓝色短裙已经在衣柜里', output)
|
||||
for phrase in ('我记', '这话', '我认', '行，', 'internal-id', '不复述的完整细节', '{'):
|
||||
self.assertNotIn(phrase, output)
|
||||
|
||||
def test_polishing_summary_describes_saved_outcome_not_relationship(self):
|
||||
output = agent_loop._polish_facts([{'tool':'wardrobe', 'action':'add', 'ok':True,
|
||||
'data':{'item':{'id':'internal', 'name':'蓝色短裙'}}}])
|
||||
self.assertIn('保存衣物到衣柜', output)
|
||||
self.assertNotIn('记住衣物', output)
|
||||
|
||||
async def test_tool_receipt_is_factual_not_a_memory_performance(self):
|
||||
item = {'id':'123456789abc', 'name':'蓝色短裙', 'category':'短裙', 'description':'白色环绕条纹'}
|
||||
store = SimpleNamespace(add=lambda *args: item)
|
||||
executor = ToolExecutor(artifacts=MagicMock(), modules={'wardrobe':SimpleNamespace(wardrobe=store)})
|
||||
ctx = ToolContext(SimpleNamespace(name='tg'), 'private', 42, 42)
|
||||
result = await executor.execute('wardrobe', {'action':'add','params':{'name':'蓝色短裙'}}, ctx)
|
||||
self.assertTrue(result.ok)
|
||||
self.assertIn('已保存到衣柜', result.safe_message)
|
||||
self.assertNotIn('记住', result.safe_message)
|
||||
self.assertEqual(result.data['item']['description'], '白色环绕条纹')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Bilibili OpenAPI command tests; no real network or message delivery."""
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import bilibili
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status=200, content_type="video/mp4", chunks=(b"video",), text=""):
|
||||
self.status_code = status
|
||||
self.headers = {"Content-Type": content_type}
|
||||
self._chunks = chunks
|
||||
self.text = text
|
||||
self.closed = False
|
||||
|
||||
def iter_content(self, size):
|
||||
yield from self._chunks
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeChannel:
|
||||
name = "tg"
|
||||
|
||||
@staticmethod
|
||||
def is_admin(uid):
|
||||
return False
|
||||
|
||||
|
||||
class BilibiliTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_sanitize_bv(self):
|
||||
self.assertEqual(bilibili.sanitize_bv("bv1uT4y1P7CX"), "BV1uT4y1P7CX")
|
||||
self.assertEqual(bilibili.sanitize_bv(" BV1uT4y1P7CX "), "BV1uT4y1P7CX")
|
||||
self.assertIsNone(bilibili.sanitize_bv("1uT4y1P7CX"))
|
||||
self.assertIsNone(bilibili.sanitize_bv("https://www.bilibili.com/video/BV1uT4y1P7CX"))
|
||||
self.assertIsNone(bilibili.sanitize_bv("BV1uT4y1P7CX extra"))
|
||||
|
||||
async def test_handle_sends_one_file_and_cleans_temp(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
video = Path(directory) / "video.mp4"
|
||||
video.write_bytes(b"video")
|
||||
sender = AsyncMock(return_value=True)
|
||||
with patch("channel.send_message", new=sender), \
|
||||
patch.object(bilibili, "download_video", return_value=video):
|
||||
await bilibili.handle_bilibili(
|
||||
FakeChannel(), "private", 10001, 1, ".bilibili BV1uT4y1P7CX"
|
||||
)
|
||||
self.assertEqual(sender.await_count, 2)
|
||||
segments = sender.await_args.args[3]
|
||||
self.assertEqual(len(segments), 1)
|
||||
self.assertEqual(segments[0]["type"], "file")
|
||||
self.assertTrue(segments[0]["data"]["file"].startswith("file://"))
|
||||
self.assertIn("bv=BV1uT4y1P7CX", segments[0]["data"]["manual"])
|
||||
self.assertFalse(video.exists())
|
||||
|
||||
async def test_invalid_input_does_not_call_openapi(self):
|
||||
sender = AsyncMock(return_value=True)
|
||||
with patch("channel.send_message", new=sender), \
|
||||
patch.object(bilibili, "download_video") as download:
|
||||
await bilibili.handle_bilibili(FakeChannel(), "private", 10001, 1, ".bilibili nope")
|
||||
download.assert_not_called()
|
||||
self.assertIn("格式", sender.await_args.args[3])
|
||||
|
||||
async def test_openapi_error_is_reported_and_temp_is_cleaned(self):
|
||||
sender = AsyncMock(return_value=True)
|
||||
with patch("channel.send_message", new=sender), \
|
||||
patch.object(bilibili, "download_video", side_effect=RuntimeError("OpenAPI HTTP 502")):
|
||||
await bilibili.handle_bilibili(
|
||||
FakeChannel(), "private", 10001, 1, ".bilibili BV1uT4y1P7CX"
|
||||
)
|
||||
self.assertIn("OpenAPI HTTP 502", sender.await_args.args[3])
|
||||
|
||||
def test_download_writes_stream_and_closes_response(self):
|
||||
response = FakeResponse(chunks=(b"abc", b"def"))
|
||||
with tempfile.TemporaryDirectory() as directory, \
|
||||
patch.object(bilibili, "CACHE_DIR", Path(directory)), \
|
||||
patch("requests.get", return_value=response) as get:
|
||||
path = bilibili.download_video("BV1uT4y1P7CX")
|
||||
try:
|
||||
self.assertEqual(path.read_bytes(), b"abcdef")
|
||||
self.assertEqual(path.suffix, ".mp4")
|
||||
get.assert_called_once_with(
|
||||
bilibili.OPENAPI_URL,
|
||||
params={"bv": "BV1uT4y1P7CX"},
|
||||
stream=True,
|
||||
timeout=bilibili.HTTP_TIMEOUT,
|
||||
)
|
||||
self.assertTrue(response.closed)
|
||||
finally:
|
||||
bilibili.cleanup_file(path)
|
||||
|
||||
def test_download_http_error_includes_detail_and_closes_response(self):
|
||||
response = FakeResponse(
|
||||
status=502,
|
||||
content_type="application/json",
|
||||
chunks=(),
|
||||
text='{"error":"download failed"}',
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as directory, \
|
||||
patch.object(bilibili, "CACHE_DIR", Path(directory)), \
|
||||
patch("requests.get", return_value=response):
|
||||
with self.assertRaisesRegex(RuntimeError, "download failed"):
|
||||
bilibili.download_video("BV1uT4y1P7CX")
|
||||
self.assertTrue(response.closed)
|
||||
self.assertEqual(list(Path(directory).iterdir()), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,380 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
# These tests never perform HTTP. Keep them offline on minimal Python installs.
|
||||
try:
|
||||
import aiohttp # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
aiohttp_stub = types.ModuleType("aiohttp")
|
||||
|
||||
class OfflineClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise RuntimeError("HTTP is disabled in offline chat migration tests")
|
||||
|
||||
aiohttp_stub.ClientSession = OfflineClientSession
|
||||
sys.modules["aiohttp"] = aiohttp_stub
|
||||
|
||||
from chat import engine
|
||||
import migrate_chat_data
|
||||
|
||||
DEFAULT_ENGINE_PATHS = {
|
||||
"history": engine.HISTORY_DIR,
|
||||
"persona": engine.PROMPT_FILE,
|
||||
"lorebook": engine.LOREBOOK_PATH,
|
||||
}
|
||||
|
||||
|
||||
class RecordingChannel:
|
||||
def __init__(self, name="tg"):
|
||||
self.name = name
|
||||
self.sent = []
|
||||
|
||||
async def send(self, msg_type, target_id, content):
|
||||
self.sent.append((msg_type, target_id, content))
|
||||
return True
|
||||
|
||||
|
||||
class EngineTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory(prefix="chat-engine-test-")
|
||||
root = Path(self.tmp.name)
|
||||
self.paths = {
|
||||
"BASE_DIR": str(root),
|
||||
"HISTORY_DIR": str(root / "history"),
|
||||
"GROUP_HISTORY_DIR": str(root / "history" / "group"),
|
||||
"MEMORY_DIR": str(root / "memory"),
|
||||
"STATUS_DIR": str(root / "status"),
|
||||
"SETTING_FILE": str(root / "setting.json"),
|
||||
"STICKERS_DIR": str(root / "stickers"),
|
||||
"PROMPTS_DIR": str(root / "prompts"),
|
||||
"LOREBOOK_PATH": str(root / "prompts" / "lorebook.json"),
|
||||
}
|
||||
self.path_patch = mock.patch.multiple(engine, **self.paths)
|
||||
self.path_patch.start()
|
||||
for key in ("HISTORY_DIR", "GROUP_HISTORY_DIR", "MEMORY_DIR", "STATUS_DIR", "STICKERS_DIR", "PROMPTS_DIR"):
|
||||
os.makedirs(self.paths[key], exist_ok=True)
|
||||
engine._delivery_targets.clear()
|
||||
engine._user_channel.clear()
|
||||
engine._history_revisions.clear()
|
||||
engine._websocket = None
|
||||
|
||||
def tearDown(self):
|
||||
engine._delivery_targets.clear()
|
||||
engine._user_channel.clear()
|
||||
engine._history_revisions.clear()
|
||||
engine._websocket = None
|
||||
self.path_patch.stop()
|
||||
self.tmp.cleanup()
|
||||
|
||||
def test_package_paths_and_prompt_defaults(self):
|
||||
package = Path(engine.__file__).resolve().parent
|
||||
self.assertEqual(Path(DEFAULT_ENGINE_PATHS["history"]), package / "history")
|
||||
self.assertEqual(Path(DEFAULT_ENGINE_PATHS["persona"]), package.parent / "prompts" / "persona.txt")
|
||||
self.assertEqual(Path(DEFAULT_ENGINE_PATHS["lorebook"]), package.parent / "prompts" / "lorebook.json")
|
||||
# The temporary prompt directory is empty: built-in defaults must work.
|
||||
self.assertIn("JSON", engine.load_instruction("agent_instructions.txt"))
|
||||
self.assertIn("群聊", engine.load_instruction("group_instructions.txt"))
|
||||
Path(engine.PROMPTS_DIR, "compact_instructions.txt").write_text("自定义压缩指令", encoding="utf-8")
|
||||
self.assertEqual(engine.load_instruction("compact_instructions.txt"), "自定义压缩指令")
|
||||
|
||||
def test_private_model_request_has_no_legacy_debounce(self):
|
||||
self.assertEqual(engine.USER_IDLE_WAIT_MIN, 0)
|
||||
self.assertEqual(engine.USER_IDLE_WAIT_MAX, 0)
|
||||
|
||||
def test_explicit_delivery_never_uses_uid_as_target(self):
|
||||
uid = 9000000001
|
||||
channel = RecordingChannel()
|
||||
engine.set_delivery(uid, channel, 10000000042)
|
||||
ok = asyncio.run(engine.send_private_message(None, uid, "你好。"))
|
||||
self.assertTrue(ok)
|
||||
self.assertEqual(channel.sent, [("private", 10000000042, "你好")])
|
||||
self.assertNotEqual(channel.sent[0][1], uid)
|
||||
|
||||
def test_restart_recovers_one_identity_target_without_mirror_fanout(self):
|
||||
uid = 9000000001
|
||||
channel = RecordingChannel()
|
||||
engine.set_websocket(channel) # only channel registry fallback; no delivery binding survives
|
||||
with mock.patch("identity.push_target", return_value=("tg", 10000000042)) as preferred, \
|
||||
mock.patch("identity.push_targets", return_value=[("tg", 10000000042), ("qq", 12345)]) as fanout:
|
||||
ok = asyncio.run(engine.send_private_message(None, uid, "单条投递"))
|
||||
self.assertTrue(ok)
|
||||
preferred.assert_called_once_with(uid)
|
||||
fanout.assert_not_called()
|
||||
self.assertEqual(channel.sent, [("private", 10000000042, "单条投递")])
|
||||
|
||||
def test_qq_only_identity_uses_first_target_only(self):
|
||||
uid = 9000000001
|
||||
channel = RecordingChannel("qq")
|
||||
engine.set_websocket(channel)
|
||||
with mock.patch("identity.push_target", return_value=None), \
|
||||
mock.patch("identity.push_targets", return_value=[("qq", 12345), ("qq", 67890)]):
|
||||
asyncio.run(engine.send_private_message(None, uid, "QQ"))
|
||||
self.assertEqual(channel.sent, [("private", 12345, "QQ")])
|
||||
|
||||
def test_unmapped_uid_fails_closed(self):
|
||||
channel = RecordingChannel()
|
||||
with mock.patch("identity.push_target", return_value=None), \
|
||||
mock.patch("identity.push_targets", return_value=[]):
|
||||
ok = asyncio.run(engine.send_private_message(channel, 9000000001, "不应发送"))
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(channel.sent, [])
|
||||
|
||||
def test_compact_cas_does_not_drop_concurrent_append(self):
|
||||
uid = 9000000001
|
||||
path = Path(engine.HISTORY_DIR) / (str(uid) + ".txt")
|
||||
original = ["[2026-01-01 00:00:0%d][user] line-%d" % (i, i) for i in range(5)]
|
||||
path.write_text("\n".join(original) + "\n", encoding="utf-8")
|
||||
|
||||
async def model_with_append(*args, **kwargs):
|
||||
engine.append_history(uid, "user", "concurrent-new-line")
|
||||
return "摘要"
|
||||
|
||||
with mock.patch.object(engine, "call_deepseek", side_effect=model_with_append):
|
||||
result = asyncio.run(engine.compact_history(uid, keep_recent=2))
|
||||
self.assertIn("未改写历史", result)
|
||||
lines = path.read_text(encoding="utf-8").splitlines()
|
||||
self.assertEqual(lines[:5], original)
|
||||
self.assertIn("concurrent-new-line", lines[-1])
|
||||
self.assertEqual(len(lines), 6)
|
||||
|
||||
async def model(*args, **kwargs):
|
||||
return "摘要"
|
||||
|
||||
with mock.patch.object(engine, "call_deepseek", side_effect=model):
|
||||
result = asyncio.run(engine.compact_history(uid, keep_recent=2))
|
||||
self.assertIn("已完成", result)
|
||||
lines = path.read_text(encoding="utf-8").splitlines()
|
||||
self.assertIn("[summary]", lines[0])
|
||||
self.assertEqual(len(lines), 3)
|
||||
self.assertIn("concurrent-new-line", lines[-1])
|
||||
self.assertEqual(list(Path(engine.HISTORY_DIR).glob(".chat-*.tmp")), [])
|
||||
|
||||
def test_status_read_modify_write_uses_uid_lock(self):
|
||||
uid = 9000000001
|
||||
barrier = threading.Barrier(9)
|
||||
|
||||
def worker(index):
|
||||
barrier.wait()
|
||||
if index % 2:
|
||||
engine.update_user_status(uid, pause=True)
|
||||
else:
|
||||
engine.update_user_status(uid, balance_notice_date="2026-10-29")
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
barrier.wait()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
status = engine.load_user_status(uid)
|
||||
self.assertTrue(status["pause"])
|
||||
self.assertEqual(status["balance_notice_date"], "2026-10-29")
|
||||
|
||||
def test_modelscope_counter_does_not_lose_whole_file_updates(self):
|
||||
workers = [threading.Thread(target=engine._modelscope_mark_success) for _ in range(12)]
|
||||
for thread in workers:
|
||||
thread.start()
|
||||
for thread in workers:
|
||||
thread.join()
|
||||
self.assertEqual(engine.load_setting()["modelscope_count"], 12)
|
||||
|
||||
def test_group_prompt_and_reply_never_read_private_memory(self):
|
||||
response = json.dumps({"message": [{"content": "群回复"}]}, ensure_ascii=False)
|
||||
with mock.patch.object(engine, "call_deepseek", new=mock.AsyncMock(return_value=response)) as model, \
|
||||
mock.patch.object(engine, "send_group_message", new=mock.AsyncMock()), \
|
||||
mock.patch.object(engine.memory_store, "migrate_legacy") as migrate, \
|
||||
mock.patch.object(engine.memory_store, "recall_facts") as recall, \
|
||||
mock.patch.object(engine.memory_store, "get_relation") as relation, \
|
||||
mock.patch.object(engine.memory_store, "list_entities") as entities:
|
||||
asyncio.run(engine._call_group_model_and_reply_impl(None, 1000000000999, "大家好"))
|
||||
migrate.assert_not_called()
|
||||
recall.assert_not_called()
|
||||
relation.assert_not_called()
|
||||
entities.assert_not_called()
|
||||
prompt = model.await_args.args[0]
|
||||
self.assertIn("不要读取、推断或泄露私聊", prompt)
|
||||
|
||||
|
||||
class MigrationTestCase(unittest.TestCase):
|
||||
UID = 9000000001
|
||||
TG_NS = 10000000042
|
||||
QQ = 12345
|
||||
GROUP = 1000000000999
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory(prefix="chat-migration-test-")
|
||||
root = Path(self.tmp.name)
|
||||
self.source = root / "old-chat"
|
||||
self.accounts = root / "accounts"
|
||||
self.dest = root / "new-chat"
|
||||
self.report = root / "report.json"
|
||||
for directory in (self.source, self.accounts / "users"):
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
record = {
|
||||
"uid": self.UID,
|
||||
"tg": {"raw": 42, "ns": self.TG_NS},
|
||||
"qq": [{"id": self.QQ}],
|
||||
}
|
||||
(self.accounts / "users" / (str(self.UID) + ".json")).write_text(
|
||||
json.dumps(record), encoding="utf-8")
|
||||
for category in ("history/group", "memory", "status", "stickers"):
|
||||
(self.source / category).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
(self.source / "history" / (str(self.TG_NS) + ".txt")).write_text(
|
||||
"shared\ntg-only\n", encoding="utf-8")
|
||||
(self.source / "history" / (str(self.QQ) + ".txt")).write_text(
|
||||
"shared\nqq-only\n", encoding="utf-8")
|
||||
(self.source / "history" / "777.txt").write_text("orphan\n", encoding="utf-8")
|
||||
(self.source / "history" / "group" / (str(self.GROUP) + ".txt")).write_text(
|
||||
"group-line\n", encoding="utf-8")
|
||||
(self.source / "memory" / (str(self.TG_NS) + ".txt")).write_text("tg-memory\n", encoding="utf-8")
|
||||
(self.source / "memory" / (str(self.QQ) + ".txt")).write_text("qq-memory\n", encoding="utf-8")
|
||||
(self.source / "status" / (str(self.TG_NS) + ".json")).write_text(json.dumps({
|
||||
"pause": True, "nested": {"tg": 1},
|
||||
}), encoding="utf-8")
|
||||
(self.source / "status" / (str(self.QQ) + ".json")).write_text(json.dumps({
|
||||
"sleep": True, "nested": {"qq": 2},
|
||||
}), encoding="utf-8")
|
||||
(self.source / "status" / "777.json").write_text("{}", encoding="utf-8")
|
||||
(self.source / "setting.json").write_text(json.dumps({"chat": True}), encoding="utf-8")
|
||||
(self.source / "stickers" / "ok.webp").write_bytes(b"sticker")
|
||||
self._make_source_db()
|
||||
self.before = self._tree_hashes(self.source)
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
@staticmethod
|
||||
def _tree_hashes(root):
|
||||
return {
|
||||
str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
for path in sorted(root.rglob("*")) if path.is_file()
|
||||
}
|
||||
|
||||
def _make_source_db(self):
|
||||
con = sqlite3.connect(str(self.source / "memory.db"))
|
||||
con.executescript(migrate_chat_data.SCHEMA)
|
||||
for old_uid, content in ((self.TG_NS, "shared fact"), (self.QQ, "shared fact"),
|
||||
(self.QQ, "qq fact"), (777, "orphan fact")):
|
||||
con.execute(
|
||||
"INSERT INTO facts(uid,content,importance,keywords,created_at,archived,source)"
|
||||
" VALUES(?,?,3,'',1,0,'test')", (old_uid, content))
|
||||
for old_uid, obj in ((self.TG_NS, "cat"), (self.QQ, "dog"), (777, "orphan")):
|
||||
con.execute(
|
||||
"INSERT INTO triples(uid,subject,predicate,object,importance,valid_from,valid_until,created_at,source)"
|
||||
" VALUES(?,?, 'likes', ?,3,1,NULL,1,'test')", (old_uid, "user", obj))
|
||||
con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", (
|
||||
self.TG_NS, 10, 2, 50, 20, json.dumps([{"ts": 1, "text": "tg"}]),
|
||||
))
|
||||
con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", (
|
||||
self.QQ, 5, 3, 60, 30, json.dumps([{"ts": 2, "text": "qq"}]),
|
||||
))
|
||||
con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", (777, 1, 9, 1, 1, "[]"))
|
||||
con.execute("INSERT INTO meta VALUES('imported:%s','1')" % self.TG_NS)
|
||||
con.execute("INSERT INTO meta VALUES('custom','value')")
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
def _counts(self):
|
||||
con = sqlite3.connect(str(self.dest / "memory.db"))
|
||||
result = {
|
||||
table: con.execute("SELECT COUNT(*) FROM %s" % table).fetchone()[0]
|
||||
for table in ("facts", "triples", "relations")
|
||||
}
|
||||
result["relation_count"] = con.execute(
|
||||
"SELECT msg_count FROM relations WHERE uid=?", (self.UID,)).fetchone()[0]
|
||||
result["first_met"] = con.execute(
|
||||
"SELECT first_met FROM relations WHERE uid=?", (self.UID,)).fetchone()[0]
|
||||
result["fts"] = con.execute("SELECT COUNT(*) FROM facts_fts").fetchone()[0]
|
||||
result["fts_hits"] = con.execute(
|
||||
"SELECT COUNT(*) FROM facts_fts WHERE facts_fts MATCH 'shared'").fetchone()[0]
|
||||
con.close()
|
||||
return result
|
||||
|
||||
def test_full_migration_quarantine_and_idempotency(self):
|
||||
report = migrate_chat_data.migrate(
|
||||
self.source, self.accounts, self.dest, self.report)
|
||||
self.assertTrue(self.report.is_file())
|
||||
self.assertEqual(report["mapping"]["tg_ns_to_uid"], {str(self.TG_NS): self.UID})
|
||||
self.assertEqual(report["mapping"]["old_key_to_uid"][str(self.QQ)], self.UID)
|
||||
history = (self.dest / "history" / (str(self.UID) + ".txt")).read_text(encoding="utf-8")
|
||||
self.assertEqual(history.splitlines(), ["shared", "tg-only", "qq-only"])
|
||||
memory = (self.dest / "memory" / (str(self.UID) + ".txt")).read_text(encoding="utf-8")
|
||||
self.assertEqual(set(memory.splitlines()), {"tg-memory", "qq-memory"})
|
||||
status = json.loads((self.dest / "status" / (str(self.UID) + ".json")).read_text(encoding="utf-8"))
|
||||
self.assertEqual(status, {"pause": True, "sleep": True, "nested": {"qq": 2, "tg": 1}})
|
||||
self.assertEqual((self.dest / "history" / "group" / (str(self.GROUP) + ".txt")).read_text(
|
||||
encoding="utf-8"), "group-line\n")
|
||||
self.assertFalse((self.dest / "history" / "777.txt").exists())
|
||||
self.assertTrue((self.dest / "quarantine" / "history" / "777.txt").is_file())
|
||||
self.assertTrue((self.dest / "quarantine" / "status" / "777.json").is_file())
|
||||
self.assertEqual((self.dest / "stickers" / "ok.webp").read_bytes(), b"sticker")
|
||||
self.assertTrue(report["sqlite"]["fts_validated"])
|
||||
self.assertEqual(report["sqlite"]["unmatched"], {"facts": 1, "triples": 1, "relations": 1})
|
||||
self.assertEqual(self._counts(), {
|
||||
"facts": 2, "triples": 2, "relations": 1, "relation_count": 5,
|
||||
"first_met": 5, "fts": 2, "fts_hits": 1,
|
||||
})
|
||||
self.assertEqual(self._tree_hashes(self.source), self.before, "source was modified")
|
||||
|
||||
# A second run must not duplicate lines/facts/triples or add relation counts again.
|
||||
migrate_chat_data.migrate(self.source, self.accounts, self.dest, self.report)
|
||||
self.assertEqual(self._counts(), {
|
||||
"facts": 2, "triples": 2, "relations": 1, "relation_count": 5,
|
||||
"first_met": 5, "fts": 2, "fts_hits": 1,
|
||||
})
|
||||
self.assertEqual((self.dest / "history" / (str(self.UID) + ".txt")).read_text(
|
||||
encoding="utf-8"), history)
|
||||
self.assertEqual(self._tree_hashes(self.source), self.before, "source changed on rerun")
|
||||
|
||||
def test_existing_destination_values_win_json_conflicts(self):
|
||||
(self.dest / "status").mkdir(parents=True)
|
||||
(self.dest / "status" / (str(self.UID) + ".json")).write_text(json.dumps({
|
||||
"pause": False, "custom": "new",
|
||||
}), encoding="utf-8")
|
||||
report = migrate_chat_data.migrate(self.source, self.accounts, self.dest, self.report)
|
||||
status = json.loads((self.dest / "status" / (str(self.UID) + ".json")).read_text(encoding="utf-8"))
|
||||
self.assertFalse(status["pause"])
|
||||
self.assertEqual(status["custom"], "new")
|
||||
self.assertTrue(status["sleep"])
|
||||
self.assertTrue(any(item["kind"] == "json_value" for item in report["conflicts"]))
|
||||
|
||||
|
||||
class RawTelegramMigrationTest(unittest.TestCase):
|
||||
def test_explicit_raw_tg_provisioning_avoids_qq_guess(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
source, accounts, dest = root / "source", root / "accounts", root / "dest"
|
||||
(source / "history").mkdir(parents=True)
|
||||
accounts.mkdir(parents=True)
|
||||
(source / "history" / "123.txt").write_text("legacy\n", encoding="utf-8")
|
||||
con = sqlite3.connect(str(source / "memory.db"))
|
||||
con.executescript(migrate_chat_data.SCHEMA)
|
||||
con.execute("INSERT INTO facts(uid,content,importance,keywords,created_at,archived,source) VALUES(123,'fact',3,'',1,0,'test')")
|
||||
con.commit()
|
||||
con.close()
|
||||
report = migrate_chat_data.migrate(
|
||||
source, accounts, dest, raw_tg=True, provision_tg=True)
|
||||
self.assertEqual(len(report["provisioned_tg"]), 1)
|
||||
uid = report["provisioned_tg"][0]["uid"]
|
||||
self.assertEqual(report["provisioned_tg"][0]["tg_ns"], 10000000123)
|
||||
self.assertTrue((dest / "history" / (str(uid) + ".txt")).is_file())
|
||||
self.assertEqual(report["sqlite"]["unmatched"]["facts"], 0)
|
||||
self.assertTrue((accounts / "users" / (str(uid) + ".json")).is_file())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Regression tests for chat API, permissions, concurrency and cleanup."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from chat import engine
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status, payload=None):
|
||||
self.status = status
|
||||
self.payload = payload or {}
|
||||
async def json(self):
|
||||
return self.payload
|
||||
async def text(self):
|
||||
return json.dumps(self.payload)
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, responses):
|
||||
self.responses = responses
|
||||
self.payloads = []
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
def post(self, url, **kwargs):
|
||||
self.payloads.append(kwargs["json"])
|
||||
response = self.responses.pop(0)
|
||||
if isinstance(response, Exception):
|
||||
raise response
|
||||
return response
|
||||
|
||||
|
||||
class ChatRegressionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_official_retry_preserves_multimodal_and_reasoning(self):
|
||||
responses = [
|
||||
FakeResponse(500, {"error": "temporary"}),
|
||||
FakeResponse(200, {"choices": [{"message": {"content": "ok"}}]}),
|
||||
]
|
||||
sessions = []
|
||||
|
||||
def factory():
|
||||
session = FakeSession(responses)
|
||||
sessions.append(session)
|
||||
return session
|
||||
|
||||
content = [{"type": "text", "text": "task"}]
|
||||
with patch.object(engine.aiohttp, "ClientSession", side_effect=factory), \
|
||||
patch.object(engine.asyncio, "sleep", new=AsyncMock()):
|
||||
result = await engine._call_deepseek_official(
|
||||
"system", content, 0, presence_penalty=0.3, reasoning=True)
|
||||
self.assertEqual(result, "ok")
|
||||
self.assertEqual(len(sessions), 2)
|
||||
for session in sessions:
|
||||
self.assertEqual(session.payloads[0]["messages"][1]["content"], content)
|
||||
self.assertEqual(session.payloads[0]["thinking"], {"type": "enabled"})
|
||||
self.assertEqual(session.payloads[0]["presence_penalty"], 0.3)
|
||||
|
||||
async def test_json_structure_is_validated(self):
|
||||
valid = '{"message":[{"content":"ok"}]}'
|
||||
self.assertIsNotNone(engine.parse_json_response(valid))
|
||||
for raw in ('{"message":["x"]}', '{"message":[null]}',
|
||||
'{"message":[{}]}', '{"message":[{"content":1}]}'):
|
||||
self.assertIsNone(engine.parse_json_response(raw))
|
||||
|
||||
async def test_mixed_group_batch_disables_actor_tools(self):
|
||||
engine._group_pending_messages[1] = ["普通用户请求", "管理员插话"]
|
||||
engine._group_pending_users[1] = [2, 1]
|
||||
engine._group_pending_images[1] = []
|
||||
reply = AsyncMock()
|
||||
with patch.object(engine, "_call_group_model_and_reply", new=reply), \
|
||||
patch.object(engine.random, "uniform", return_value=0):
|
||||
await engine._process_group_messages(1)
|
||||
self.assertEqual(reply.await_args.args[3], 0)
|
||||
|
||||
async def test_private_turns_are_serialized(self):
|
||||
active = 0
|
||||
maximum = 0
|
||||
|
||||
async def impl(*args, **kwargs):
|
||||
nonlocal active, maximum
|
||||
active += 1
|
||||
maximum = max(maximum, active)
|
||||
await asyncio.sleep(0.02)
|
||||
active -= 1
|
||||
|
||||
center = type("Center", (), {
|
||||
"begin_conversation": lambda self, uid: None,
|
||||
"end_conversation": lambda self, uid: None,
|
||||
"deliver_due": AsyncMock(),
|
||||
})()
|
||||
import notifications
|
||||
with patch.object(engine, "_resolve_delivery", return_value=None), \
|
||||
patch.object(engine, "_call_model_and_reply_impl", new=impl), \
|
||||
patch.object(notifications, "center", center):
|
||||
await asyncio.gather(*[
|
||||
engine._call_model_and_reply(None, 1, f"message-{i}", False)
|
||||
for i in range(3)
|
||||
])
|
||||
self.assertEqual(maximum, 1)
|
||||
|
||||
async def test_reasoning_json_gets_one_repair_retry(self):
|
||||
valid = '{"message":[{"content":"ok"}],"tools":[{"name":"web_search","args":{"query":"x"}}]}'
|
||||
model = AsyncMock(side_effect=["not json", valid])
|
||||
with patch.object(engine, "call_deepseek", new=model):
|
||||
parsed = await engine._reasoning_json_response("system", "user")
|
||||
self.assertEqual(parsed["tools"][0]["name"], "web_search")
|
||||
self.assertEqual(model.await_count, 2)
|
||||
self.assertTrue(all(call.kwargs["reasoning"] for call in model.await_args_list))
|
||||
|
||||
async def test_reasoning_json_stops_after_two_invalid_responses(self):
|
||||
model = AsyncMock(side_effect=["bad", "still bad"])
|
||||
with patch.object(engine, "call_deepseek", new=model):
|
||||
parsed = await engine._reasoning_json_response("system", "user")
|
||||
self.assertIsNone(parsed)
|
||||
self.assertEqual(model.await_count, 2)
|
||||
|
||||
async def test_send_segments_reports_delivery_failure(self):
|
||||
with patch.object(engine, "send_private_message", new=AsyncMock(side_effect=[True, False])):
|
||||
ok = await engine.send_message_segments(None, 1, [{"content": "a"}, {"content": "b"}])
|
||||
self.assertFalse(ok)
|
||||
|
||||
async def test_pending_images_are_cleaned_after_consumption(self):
|
||||
with tempfile.NamedTemporaryFile(delete=False) as handle:
|
||||
path = Path(handle.name)
|
||||
engine._user_pending_messages[1] = ["读图"]
|
||||
engine._user_pending_images[1] = [{"url": str(path)}]
|
||||
with patch.object(engine, "_call_model_and_reply", new=AsyncMock()):
|
||||
await engine._process_pending_messages(1, 0)
|
||||
self.assertFalse(path.exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Credential assignments stay configurable without importing live services."""
|
||||
import ast
|
||||
import os
|
||||
from pathlib import Path
|
||||
import secrets
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MAPPINGS = (
|
||||
('bot-utility/for_admin_komari_agent.py', 'KOMARI_PASSWORD', 'KOMARI_PASSWORD'),
|
||||
('bot-utility/pixiv.py', 'CF_PASSWORD', 'PIXIV_CF_PASSWORD'),
|
||||
('mail-bridge/mailWeb.py', 'HMAIL_ADMIN_PASSWORD', 'HMAIL_ADMIN_PASSWORD'),
|
||||
('mail-bridge/mailWeb.py', 'REGISTER_AUTH_CODE', 'MAIL_REGISTER_AUTH_CODE'),
|
||||
('mail-bridge/mailWeb.py', 'app.secret_key', 'MAIL_FLASK_SECRET_KEY'),
|
||||
)
|
||||
|
||||
|
||||
def expression(relative, target):
|
||||
tree = ast.parse((ROOT / relative).read_text(encoding='utf-8'))
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.Assign) and any(ast.unparse(t) == target for t in node.targets):
|
||||
return ast.Expression(node.value)
|
||||
raise AssertionError('Credential assignment missing: ' + target)
|
||||
|
||||
|
||||
class EnvironmentCredentialTests(unittest.TestCase):
|
||||
def test_each_credential_comes_from_its_environment_variable(self):
|
||||
for path, target, env in MAPPINGS:
|
||||
with self.subTest(target=target), patch.dict(os.environ, {env: 'temporary-unit-test-value'}):
|
||||
value = eval(compile(expression(path, target), '<credential-assignment>', 'eval'),
|
||||
{'os': os, 'secrets': secrets})
|
||||
self.assertEqual(value, 'temporary-unit-test-value')
|
||||
|
||||
def test_missing_passwords_default_empty_not_embedded_credentials(self):
|
||||
for path, target, _ in MAPPINGS[:-1]:
|
||||
with self.subTest(target=target), patch.dict(os.environ, {}, clear=True):
|
||||
value = eval(compile(expression(path, target), '<credential-assignment>', 'eval'),
|
||||
{'os': os, 'secrets': secrets})
|
||||
self.assertEqual(value, '')
|
||||
|
||||
def test_unconfigured_flask_key_is_ephemeral_not_shared(self):
|
||||
path, target, _ = MAPPINGS[-1]
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
code = compile(expression(path, target), '<credential-assignment>', 'eval')
|
||||
one = eval(code, {'os': os, 'secrets': secrets})
|
||||
two = eval(code, {'os': os, 'secrets': secrets})
|
||||
self.assertTrue(one)
|
||||
self.assertNotEqual(one, two)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Telegram Help must retain inline buttons in both command and chat modes."""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import main
|
||||
|
||||
|
||||
class FakeTGChannel:
|
||||
tg_username = "test_bot"
|
||||
|
||||
@staticmethod
|
||||
def is_admin(uid):
|
||||
return True
|
||||
|
||||
|
||||
def private_message(text="/help"):
|
||||
return SimpleNamespace(
|
||||
from_user=SimpleNamespace(id=123, is_bot=False),
|
||||
chat=SimpleNamespace(id=123, type="private"),
|
||||
text=text,
|
||||
caption=None,
|
||||
photo=[],
|
||||
entities=[],
|
||||
)
|
||||
|
||||
|
||||
class TelegramHelpTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def run_root_help(self, mode):
|
||||
bot = SimpleNamespace(send_message=AsyncMock(), id=999)
|
||||
store = SimpleNamespace(get_mode=lambda *args: mode)
|
||||
with patch.object(main, "ModeStore", return_value=store), \
|
||||
patch.object(main.identity, "resolve", return_value=1), \
|
||||
patch.object(main.help_ui, "chat_help", new=AsyncMock(return_value="chat help")) as chat_help, \
|
||||
patch.object(main, "handle_message", new=AsyncMock()) as dispatch:
|
||||
await main.make_tg_handler(FakeTGChannel())(private_message(), bot)
|
||||
return bot, chat_help, dispatch
|
||||
|
||||
async def test_command_root_help_has_buttons(self):
|
||||
bot, chat_help, dispatch = await self.run_root_help("command")
|
||||
kwargs = bot.send_message.await_args.kwargs
|
||||
self.assertIsNotNone(kwargs.get("reply_markup"))
|
||||
self.assertIn("服务列表", bot.send_message.await_args.args[1])
|
||||
chat_help.assert_not_awaited()
|
||||
dispatch.assert_not_awaited()
|
||||
|
||||
async def test_chat_root_help_has_buttons_and_skips_legacy_dispatch(self):
|
||||
bot, chat_help, dispatch = await self.run_root_help("chat")
|
||||
kwargs = bot.send_message.await_args.kwargs
|
||||
self.assertIsNotNone(kwargs.get("reply_markup"))
|
||||
self.assertEqual(bot.send_message.await_args.args[1], "chat help")
|
||||
chat_help.assert_awaited_once_with(1, None, True)
|
||||
dispatch.assert_not_awaited()
|
||||
|
||||
async def test_chat_help_callback_keeps_buttons_and_mode_aware_content(self):
|
||||
bot = SimpleNamespace(send_message=AsyncMock())
|
||||
callback = SimpleNamespace(
|
||||
from_user=SimpleNamespace(id=123),
|
||||
message=SimpleNamespace(chat=SimpleNamespace(id=123, type="private")),
|
||||
data="help:1",
|
||||
answer=AsyncMock(),
|
||||
)
|
||||
store = SimpleNamespace(get_mode=lambda *args: "chat")
|
||||
with patch.object(main, "ModeStore", return_value=store), \
|
||||
patch.object(main.identity, "resolve", return_value=1), \
|
||||
patch.object(main.help_ui, "chat_help", new=AsyncMock(return_value="setu chat help")) as chat_help:
|
||||
await main.make_tg_help_callback(FakeTGChannel())(callback, bot)
|
||||
chat_help.assert_awaited_once_with(1, "setu", True)
|
||||
self.assertEqual(bot.send_message.await_args.args[1], "setu chat help")
|
||||
self.assertIsNotNone(bot.send_message.await_args.kwargs.get("reply_markup"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Komari notifications and natural-style prompt integration tests."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import for_admin_komari_agent as komari
|
||||
import help_ui
|
||||
import notification_render
|
||||
import notifications
|
||||
from chat import engine
|
||||
|
||||
|
||||
class KomariNotificationTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_komari_alert_enters_structured_notification_pipeline(self):
|
||||
with patch("identity.admin_uids", return_value=[1]), \
|
||||
patch("identity.push_targets", return_value=[("tg", 10001)]), \
|
||||
patch.object(notifications.center, "publish", return_value="event-1") as publish:
|
||||
accepted = await komari._notify_admins(
|
||||
"[Komari] 服务器「Ama-10」当前处于离线状态。",
|
||||
node_uuid="node-1",
|
||||
server="Ama-10",
|
||||
status="offline",
|
||||
)
|
||||
self.assertTrue(accepted)
|
||||
event = publish.call_args.args[0]
|
||||
self.assertEqual(event.kind, "komari")
|
||||
self.assertEqual(event.uid, 1)
|
||||
self.assertEqual(event.payload["server"], "Ama-10")
|
||||
self.assertEqual(event.payload["status"], "offline")
|
||||
self.assertIn("当前处于离线状态", event.legacy_text)
|
||||
|
||||
async def test_komari_alert_without_targets_does_not_advance(self):
|
||||
with patch("identity.admin_uids", return_value=[1]), \
|
||||
patch("identity.push_targets", return_value=[]), \
|
||||
patch.object(notifications.center, "publish") as publish:
|
||||
accepted = await komari._notify_admins(
|
||||
"offline", node_uuid="node-1", server="Ama-10", status="offline"
|
||||
)
|
||||
self.assertFalse(accepted)
|
||||
publish.assert_not_called()
|
||||
|
||||
async def test_chat_renderer_can_naturalize_komari_event(self):
|
||||
event = notifications.Event(
|
||||
kind="komari",
|
||||
dedupe_key="komari:test",
|
||||
legacy_text="[Komari] 服务器「Ama-10」当前处于离线状态。",
|
||||
uid=1,
|
||||
payload={"server": "Ama-10", "status": "offline", "source": "Komari"},
|
||||
)
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value="Ama-10 掉线了。")) as model:
|
||||
rendered = await notification_render.render_event(event, "chat", 10001)
|
||||
self.assertEqual(rendered, "Ama-10 掉线了。")
|
||||
self.assertIn("自然表达补充", model.await_args.args[0])
|
||||
|
||||
|
||||
class StylePromptTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_base_prompt_includes_natural_style_supplement(self):
|
||||
self.assertIn("自然表达补充", engine.BASE_PROMPT)
|
||||
self.assertIn("话不算多，先把", engine.BASE_PROMPT)
|
||||
self.assertIn("我先跟你说一声", engine.BASE_PROMPT)
|
||||
self.assertIn("不过还不能", engine.BASE_PROMPT)
|
||||
self.assertIn("已知且与当前回答直接相关的客观限制", engine.BASE_PROMPT)
|
||||
self.assertIn("不影响 JSON", engine.BASE_PROMPT)
|
||||
|
||||
def test_active_prompt_has_varied_events_without_waiting_for_user(self):
|
||||
prompt = engine.load_instruction("notification_instructions.txt")
|
||||
self.assertIn("随机发生的生活感", prompt)
|
||||
self.assertIn("不是从不出门", prompt)
|
||||
self.assertIn("生活只围着用户转", prompt)
|
||||
self.assertIn("不必每条都关心用户", prompt)
|
||||
|
||||
async def test_help_no_longer_requires_self_introduction(self):
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value="能做的事如下")) as model:
|
||||
result = await help_ui.chat_help(1)
|
||||
self.assertEqual(result, "能做的事如下")
|
||||
prompt = model.await_args.args[0]
|
||||
self.assertIn("不必自我介绍", prompt)
|
||||
self.assertNotIn("必须先简短介绍自己", prompt)
|
||||
|
||||
async def test_help_fallback_also_avoids_self_introduction_and_preview(self):
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value="")):
|
||||
result = await help_ui.chat_help(1)
|
||||
self.assertNotIn("我是白花", result)
|
||||
self.assertNotIn("话不算多", result)
|
||||
self.assertNotIn("先理清顺序", result)
|
||||
self.assertIn("直接说你想做什么就行", result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,361 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""旧业务模块回归测试；网络、邮件、签到及生产文件路径全部隔离。"""
|
||||
import asyncio
|
||||
import queue
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
# The minimal audit container lacks aiohttp. Tests replace every network object,
|
||||
# so a tiny import stub is sufficient there; production uses the real dependency.
|
||||
try:
|
||||
import aiohttp # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
aiohttp_stub = types.ModuleType("aiohttp")
|
||||
|
||||
class ClientTimeout:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_stub.ClientTimeout = ClientTimeout
|
||||
aiohttp_stub.ClientSession = None
|
||||
sys.modules["aiohttp"] = aiohttp_stub
|
||||
|
||||
import for_admin_komari_agent as komari # noqa: E402
|
||||
import ginkgostu # noqa: E402
|
||||
import identity_cmd # noqa: E402
|
||||
import jm # noqa: E402
|
||||
import live_monitor # noqa: E402
|
||||
import mail # noqa: E402
|
||||
import migrate_chat_data # noqa: E402
|
||||
import pixiv # noqa: E402
|
||||
import setu # noqa: E402
|
||||
|
||||
|
||||
class GinkgoRegressionTests(unittest.TestCase):
|
||||
URL = "https://ginkgostu.dufe.edu.cn/course/detail/1K/batch/2/signin/9/2"
|
||||
|
||||
def test_retry_status_confirms_success_without_another_commit(self):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"data": {"signInModel": {"type": 2, "title": "t"}, "signStatus": 1}
|
||||
}
|
||||
session = MagicMock()
|
||||
with patch.object(ginkgostu, "call_api", return_value=response), \
|
||||
patch.object(ginkgostu, "commit_sign_in") as commit:
|
||||
result = ginkgostu.fetch_sign_in_info(
|
||||
"u", "p", self.URL, sess=session, token="token",
|
||||
avoid_duplicate=True)
|
||||
self.assertTrue(result["commit"]["success"])
|
||||
commit.assert_not_called()
|
||||
session.close.assert_not_called() # 调用方所有的缓存会话不能关闭
|
||||
|
||||
def test_owned_signin_session_closes_on_query_failure(self):
|
||||
session = MagicMock()
|
||||
with patch.object(ginkgostu, "login", return_value=(session, "token", "n")), \
|
||||
patch.object(ginkgostu, "call_api",
|
||||
side_effect=ginkgostu.UpstreamError("bad response")):
|
||||
with self.assertRaises(ginkgostu.UpstreamError):
|
||||
ginkgostu.fetch_sign_in_info("u", "p", self.URL)
|
||||
session.close.assert_called_once_with()
|
||||
|
||||
def test_login_failure_closes_sso_session(self):
|
||||
session = MagicMock()
|
||||
with patch.object(ginkgostu, "new_session", return_value=session), \
|
||||
patch.object(ginkgostu, "get_unify_config",
|
||||
side_effect=ginkgostu.UpstreamError("failed")):
|
||||
with self.assertRaises(ginkgostu.UpstreamError):
|
||||
ginkgostu.login("u", "p")
|
||||
session.close.assert_called_once_with()
|
||||
|
||||
def test_refresh_closes_replaced_cached_session(self):
|
||||
old = MagicMock()
|
||||
new = MagicMock()
|
||||
cached = ginkgostu.GinkgoSession("u", "p")
|
||||
cached.req = old
|
||||
with patch.object(ginkgostu, "login", return_value=(new, "new-token", "n")):
|
||||
cached.refresh()
|
||||
old.close.assert_called_once_with()
|
||||
self.assertIs(cached.req, new)
|
||||
|
||||
def test_signin_api_forwards_retry_status_check(self):
|
||||
with patch.object(ginkgostu, "_call_sign_in_local",
|
||||
return_value=(True, "ok", {}, False)) as local:
|
||||
ginkgostu.call_sign_in_api("u", "p", self.URL, check_status=True)
|
||||
self.assertTrue(local.call_args.kwargs["check_status"])
|
||||
|
||||
|
||||
class MailRegressionTests(unittest.TestCase):
|
||||
def test_transport_result_unknown_is_not_a_plain_send_failure(self):
|
||||
user = {"addresses": ["sender@rrrosmontis.icu"], "default": "",
|
||||
"push": {"enabled": False, "last_seen": {}}}
|
||||
with patch.object(mail, "_load_user", return_value=user), \
|
||||
patch.object(mail, "_ComSession", side_effect=RuntimeError("lost response")):
|
||||
result = asyncio.run(mail.send_with_attachments(
|
||||
42, "sender@rrrosmontis.icu", "to@example.net", "s", "b", []))
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error_code"], "send_unknown")
|
||||
|
||||
def test_explicit_bridge_rejection_remains_send_failed(self):
|
||||
user = {"addresses": ["sender@rrrosmontis.icu"], "default": "",
|
||||
"push": {"enabled": False, "last_seen": {}}}
|
||||
app = MagicMock()
|
||||
app.post.return_value = {"accepted": False}
|
||||
with patch.object(mail, "_load_user", return_value=user), \
|
||||
patch.object(mail, "_ComSession") as session:
|
||||
session.return_value.__enter__.return_value = app
|
||||
result = asyncio.run(mail.send_with_attachments(
|
||||
42, "sender@rrrosmontis.icu", "to@example.net", "s", "b", []))
|
||||
self.assertEqual(result["error_code"], "send_failed")
|
||||
|
||||
|
||||
class LiveMonitorRegressionTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.client = live_monitor.app.test_client()
|
||||
live_monitor._live_sessions.clear()
|
||||
self.records = {"follows": {"42": [{
|
||||
"chat_type": "group", "chat_id": 7, "channel": "qq", "site": "1",
|
||||
}]}}
|
||||
self.secret = patch.object(live_monitor, "SRS_CALLBACK_SECRET", "test-secret")
|
||||
self.load = patch.object(live_monitor, "_load_records", return_value=self.records)
|
||||
self.notifications = patch.object(live_monitor, "notification_queue", queue.Queue())
|
||||
self.secret.start()
|
||||
self.load.start()
|
||||
self.notifications.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.notifications.stop()
|
||||
self.load.stop()
|
||||
self.secret.stop()
|
||||
live_monitor._live_sessions.clear()
|
||||
|
||||
def test_missing_srs_secret_fails_closed(self):
|
||||
with patch.object(live_monitor, "SRS_CALLBACK_SECRET", ""):
|
||||
response = self.client.post(
|
||||
"/remote_srs/on_publish", json={"stream": "room_42"})
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertTrue(live_monitor.notification_queue.empty())
|
||||
|
||||
def test_srs_auth_dedupe_and_unpublish_rearm(self):
|
||||
payload = {"stream": "room_42"}
|
||||
denied = self.client.post("/remote_srs/on_publish", json=payload)
|
||||
self.assertEqual(denied.status_code, 403)
|
||||
self.assertTrue(live_monitor.notification_queue.empty())
|
||||
|
||||
for _ in range(2):
|
||||
accepted = self.client.post(
|
||||
"/remote_srs/on_publish?token=test-secret", json=payload)
|
||||
self.assertEqual(accepted.status_code, 200)
|
||||
self.assertEqual(live_monitor.notification_queue.qsize(), 1)
|
||||
|
||||
denied_end = self.client.post("/remote_srs/on_unpublish", json=payload)
|
||||
self.assertEqual(denied_end.status_code, 403)
|
||||
ended = self.client.post(
|
||||
"/remote_srs/on_unpublish", json=payload,
|
||||
headers={"X-SRS-Token": "test-secret"})
|
||||
self.assertEqual(ended.status_code, 200)
|
||||
self.client.post("/remote_srs/on_publish?token=test-secret", json=payload)
|
||||
self.assertEqual(live_monitor.notification_queue.qsize(), 2)
|
||||
|
||||
|
||||
class KomariAndIdentityTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_recent_exception_preserves_state_without_offline_notice(self):
|
||||
komari._last_state = {"node-1": {"name": "one", "offline": False}}
|
||||
session = MagicMock()
|
||||
session.__aenter__ = AsyncMock(return_value=session)
|
||||
session.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
async def get_json(_session, path, token=None):
|
||||
if path == "/api/nodes":
|
||||
return {"status": "success", "data": [{"uuid": "node-1", "name": "one"}]}
|
||||
raise RuntimeError("temporary recent failure")
|
||||
|
||||
with patch.object(komari, "_login", AsyncMock(return_value="key")), \
|
||||
patch.object(komari.aiohttp, "ClientSession", return_value=session), \
|
||||
patch.object(komari, "_get_json", side_effect=get_json), \
|
||||
patch.object(komari, "_notify_admins", AsyncMock()) as notify:
|
||||
await komari._check_servers_once()
|
||||
self.assertEqual(komari._last_state["node-1"]["offline"], False)
|
||||
notify.assert_not_awaited()
|
||||
|
||||
async def test_notification_send_false_is_not_reported_as_sent(self):
|
||||
channel = object()
|
||||
with patch.object(identity_cmd, "get_channel", return_value=channel), \
|
||||
patch.object(identity_cmd, "send_message", AsyncMock(return_value=False)):
|
||||
self.assertFalse(await identity_cmd._notify("qq", 42, "text"))
|
||||
with patch.object(identity_cmd, "get_channel", return_value=channel), \
|
||||
patch.object(identity_cmd, "send_message", AsyncMock(return_value=True)):
|
||||
self.assertTrue(await identity_cmd._notify("qq", 42, "text"))
|
||||
|
||||
|
||||
class QuarantineRegressionTests(unittest.TestCase):
|
||||
def test_same_name_quarantine_files_never_overwrite(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
source = root / "source"
|
||||
dest = root / "dest"
|
||||
first = source / "a" / "same.txt"
|
||||
second = source / "b" / "same.txt"
|
||||
first.parent.mkdir(parents=True)
|
||||
second.parent.mkdir(parents=True)
|
||||
migrator = migrate_chat_data.ChatMigrator.__new__(
|
||||
migrate_chat_data.ChatMigrator)
|
||||
migrator.source = source
|
||||
migrator.dest = dest
|
||||
migrator.report = {"quarantine": []}
|
||||
|
||||
first.write_bytes(b"first")
|
||||
second.write_bytes(b"second")
|
||||
migrator._quarantine(first, "history", "one")
|
||||
migrator._quarantine(second, "history", "two")
|
||||
first.write_bytes(b"changed later")
|
||||
migrator._quarantine(first, "history", "three")
|
||||
|
||||
saved = sorted(p.read_bytes() for p in (dest / "quarantine" / "history").iterdir())
|
||||
self.assertEqual(saved, [b"changed later", b"first", b"second"])
|
||||
|
||||
|
||||
class JmRegressionTests(unittest.TestCase):
|
||||
def test_local_pack_import_and_stale_workdir_cleanup(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
album = "350234"
|
||||
stale = root / album / "old-chapter" / "stale.jpg"
|
||||
stale.parent.mkdir(parents=True)
|
||||
stale.write_bytes(b"stale")
|
||||
plan = ("title", [(101,)], 1, 1, "title.zip")
|
||||
|
||||
class Downloader:
|
||||
def __init__(self, _option):
|
||||
pass
|
||||
|
||||
def download_photo(self, photo_id):
|
||||
target = root / album / str(photo_id)
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
(target / "image.webp").write_bytes(b"not-an-image-is-copied")
|
||||
|
||||
class Options:
|
||||
@staticmethod
|
||||
def construct(_config):
|
||||
return object()
|
||||
|
||||
class Config:
|
||||
option_class = Options
|
||||
|
||||
@staticmethod
|
||||
def downloader_class():
|
||||
return Downloader
|
||||
|
||||
jmcomic_stub = types.ModuleType("jmcomic")
|
||||
jmcomic_stub.JmModuleConfig = Config
|
||||
with patch.object(jm, "JM_DIR", str(root)), \
|
||||
patch.dict(sys.modules, {"jmcomic": jmcomic_stub}):
|
||||
title, zip_path, count, ok, failed = jm._download_and_pack(
|
||||
album, 1, 1, plan=plan)
|
||||
self.assertEqual((title, count, ok, failed), ("title", 1, 1, 0))
|
||||
self.assertTrue(zip_path.is_file())
|
||||
self.assertFalse(stale.exists())
|
||||
|
||||
|
||||
class ImageCleanupTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_pixiv_cancellation_cleans_cache_and_releases_lock(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pid = "424242"
|
||||
local = Path(tmp) / (pid + ".jpg")
|
||||
local.write_bytes(b"image")
|
||||
lock = pixiv.get_pid_lock(pid)
|
||||
with patch.object(pixiv, "SEND_IMAGE", 1), \
|
||||
patch.object(pixiv, "RETAIN_MODE", "oss"), \
|
||||
patch.object(pixiv, "IS_NAPCAT_SERVER_OSS_AVAILABLE", False), \
|
||||
patch.object(pixiv, "OSS_SAVE_DIR", str(Path(tmp) / "oss")), \
|
||||
patch.object(pixiv, "_load_proxy_ports", return_value=[]), \
|
||||
patch.object(pixiv, "download_image", return_value=local), \
|
||||
patch.object(pixiv, "image_base64", return_value="base64://x"), \
|
||||
patch("channel.send_message", AsyncMock(return_value=True)), \
|
||||
patch.object(pixiv.asyncio, "sleep",
|
||||
AsyncMock(side_effect=asyncio.CancelledError)):
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await pixiv.handle_pixiv(MagicMock(name="qq"), "private", 1, 2,
|
||||
".pixiv " + pid)
|
||||
self.assertFalse(local.exists())
|
||||
self.assertFalse(lock.locked())
|
||||
|
||||
async def test_setu_cancellation_cleans_cache_and_releases_lock(self):
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
async def json(self):
|
||||
return {"data": [{"pid": 77, "urls": {
|
||||
"original": "https://i.pixiv.re/img/77.jpg"}}]}
|
||||
|
||||
class Context:
|
||||
async def __aenter__(self):
|
||||
return Response()
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
class Session:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
return Context()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
local = Path(tmp) / "77.jpg"
|
||||
local.write_bytes(b"image")
|
||||
lock = pixiv.get_pid_lock(77)
|
||||
with patch.object(setu, "SEND_IMAGE", 1), \
|
||||
patch.object(setu, "RETAIN_MODE", "oss"), \
|
||||
patch.object(setu, "IS_NAPCAT_SERVER_OSS_AVAILABLE", False), \
|
||||
patch.object(setu, "OSS_SAVE_DIR", str(Path(tmp) / "oss")), \
|
||||
patch.object(setu, "_load_proxy_ports", return_value=[]), \
|
||||
patch.object(setu, "count_oss_files", return_value=0), \
|
||||
patch.object(setu.aiohttp, "ClientSession", Session), \
|
||||
patch.object(pixiv, "download_image", return_value=local), \
|
||||
patch.object(pixiv, "image_base64", return_value="base64://x"), \
|
||||
patch("channel.send_message", AsyncMock(return_value=True)), \
|
||||
patch.object(setu.asyncio, "sleep",
|
||||
AsyncMock(side_effect=asyncio.CancelledError)):
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await setu.handle_setu([], None, None, "private", 1, MagicMock(name="qq"))
|
||||
self.assertFalse(local.exists())
|
||||
self.assertFalse(lock.locked())
|
||||
|
||||
async def test_interrupted_download_removes_partial_file(self):
|
||||
class Response:
|
||||
status_code = 200
|
||||
headers = {"Content-Type": "image/jpeg"}
|
||||
|
||||
def iter_content(self, _size):
|
||||
yield b"partial"
|
||||
raise RuntimeError("connection lost")
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
response = Response()
|
||||
with tempfile.TemporaryDirectory() as tmp, \
|
||||
patch.object(pixiv, "PIXIV_CACHE_DIR", tmp), \
|
||||
patch("requests.get", return_value=response):
|
||||
with self.assertRaises(RuntimeError):
|
||||
pixiv.download_image("88")
|
||||
self.assertEqual(list(Path(tmp).iterdir()), [])
|
||||
self.assertTrue(response.closed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,369 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""附件邮件协议、校验与 MIME 测试；全部外部邮件/COM 操作均 mock。"""
|
||||
import asyncio
|
||||
import base64
|
||||
import email
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
from flask import Flask
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE / "mail-bridge"), str(BASE)]
|
||||
|
||||
import mail # noqa: E402
|
||||
|
||||
# com_backend 在 import 时读取配置；这里只提供无真实凭据的测试配置。
|
||||
_CONFIG_TMP = tempfile.TemporaryDirectory(prefix="mail-attachment-test-config-")
|
||||
_CONFIG_PATH = Path(_CONFIG_TMP.name) / "com-private.json"
|
||||
_CONFIG_PATH.write_text(json.dumps({
|
||||
"hmail_admin_password": "not-used-by-mock-tests",
|
||||
"token": "t" * 128,
|
||||
"allowed_domains": ["rrrosmontis.icu", "aeolianchi.me"],
|
||||
}), encoding="utf-8")
|
||||
with patch.dict(os.environ, {"MAIL_COM_CONFIG": str(_CONFIG_PATH)}):
|
||||
import com_backend as backend # noqa: E402
|
||||
import admin_com # noqa: E402
|
||||
|
||||
|
||||
SENDER = "sender@rrrosmontis.icu"
|
||||
RECIPIENT = "receiver@example.net"
|
||||
AUTH = {"Authorization": "Bearer " + "t" * 128}
|
||||
ENDPOINT = "/v1/mail/api/admin/com/send_with_attachments"
|
||||
SEND_ENDPOINT = "/v1/mail/api/admin/com/send"
|
||||
|
||||
|
||||
def outgoing_payload(**overrides):
|
||||
data = {
|
||||
"from_address": SENDER,
|
||||
"to_address": RECIPIENT,
|
||||
"subject": "附件测试",
|
||||
"body_text": "正文",
|
||||
"attachments": [],
|
||||
}
|
||||
data.update(overrides)
|
||||
return data
|
||||
|
||||
|
||||
def wire_attachment(content=b"hello", filename="测试.txt",
|
||||
content_type="text/plain"):
|
||||
return {
|
||||
"filename": filename,
|
||||
"content_type": content_type,
|
||||
"data_base64": base64.b64encode(content).decode("ascii"),
|
||||
}
|
||||
|
||||
|
||||
class ClientAttachmentTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory(prefix="mail-attachment-client-")
|
||||
self.mail_dir = Path(self.tmp.name) / "mail"
|
||||
self.mail_dir.mkdir()
|
||||
self.dir_patch = patch.object(mail, "MAIL_DIR", str(self.mail_dir))
|
||||
self.dir_patch.start()
|
||||
mail._save_user(42, {
|
||||
"addresses": [SENDER],
|
||||
"default": SENDER,
|
||||
"push": {"enabled": False, "interval_minutes": 5, "last_seen": {}},
|
||||
})
|
||||
self.files = Path(self.tmp.name) / "files"
|
||||
self.files.mkdir()
|
||||
|
||||
def tearDown(self):
|
||||
self.dir_patch.stop()
|
||||
self.tmp.cleanup()
|
||||
|
||||
def send(self, paths, **kwargs):
|
||||
args = dict(user_id=42, from_address=SENDER, to_address=RECIPIENT,
|
||||
subject="主题", body_text="正文", attachment_paths=paths)
|
||||
args.update(kwargs)
|
||||
return asyncio.run(mail.send_with_attachments(**args))
|
||||
|
||||
def test_success_uses_base64_protocol_and_chinese_filename(self):
|
||||
path = self.files / "中文 报告.txt"
|
||||
content = "附件内容".encode("utf-8")
|
||||
path.write_bytes(content)
|
||||
app = MagicMock()
|
||||
app.post.return_value = {"ok": True, "accepted": True, "sent_saved": True}
|
||||
with patch.object(mail, "_ComSession") as session:
|
||||
session.return_value.__enter__.return_value = app
|
||||
result = self.send([path])
|
||||
|
||||
self.assertTrue(result["ok"])
|
||||
self.assertEqual(result["from"], SENDER)
|
||||
self.assertEqual(result["to"], RECIPIENT)
|
||||
self.assertEqual(result["attachment_names"], ["中文 报告.txt"])
|
||||
self.assertEqual(result["attachment_count"], 1)
|
||||
self.assertEqual(result["size"], len(content))
|
||||
self.assertTrue(result["sent_saved"])
|
||||
endpoint, payload = app.post.call_args.args
|
||||
self.assertEqual(endpoint, "/send_with_attachments")
|
||||
self.assertNotIn("path", payload["attachments"][0])
|
||||
self.assertEqual(payload["attachments"][0]["filename"], "中文 报告.txt")
|
||||
self.assertEqual(base64.b64decode(
|
||||
payload["attachments"][0]["data_base64"], validate=True), content)
|
||||
self.assertEqual(payload["attachments"][0]["content_type"], "text/plain")
|
||||
|
||||
def test_from_address_must_be_bound(self):
|
||||
with patch.object(mail, "_ComSession") as session:
|
||||
result = self.send([], from_address="other@rrrosmontis.icu")
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error_code"], "from_not_bound")
|
||||
session.assert_not_called()
|
||||
|
||||
def test_recipient_subject_and_body_validation(self):
|
||||
cases = [
|
||||
({"to_address": "not-an-address"}, "invalid_recipient"),
|
||||
({"subject": "bad\nsubject"}, "invalid_subject"),
|
||||
({"subject": " "}, "invalid_subject"),
|
||||
({"body_text": None}, "invalid_body"),
|
||||
]
|
||||
for overrides, code in cases:
|
||||
with self.subTest(code=code):
|
||||
with patch.object(mail, "_ComSession") as session:
|
||||
result = self.send([], **overrides)
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error_code"], code)
|
||||
session.assert_not_called()
|
||||
|
||||
def test_single_file_total_and_count_limits(self):
|
||||
one = self.files / "one.bin"
|
||||
one.write_bytes(b"12345")
|
||||
with patch.object(mail, "MAX_ATTACHMENT_BYTES", 4):
|
||||
result = self.send([one])
|
||||
self.assertEqual(result["error_code"], "file_too_large")
|
||||
|
||||
two = self.files / "two.bin"
|
||||
one.write_bytes(b"123")
|
||||
two.write_bytes(b"456")
|
||||
with patch.object(mail, "MAX_ATTACHMENT_BYTES", 4), \
|
||||
patch.object(mail, "MAX_ATTACHMENT_TOTAL_BYTES", 5):
|
||||
result = self.send([one, two])
|
||||
self.assertEqual(result["error_code"], "total_too_large")
|
||||
|
||||
with patch.object(mail, "MAX_ATTACHMENT_COUNT", 1):
|
||||
result = self.send([one, two])
|
||||
self.assertEqual(result["error_code"], "too_many_attachments")
|
||||
|
||||
def test_directory_and_special_file_are_rejected(self):
|
||||
result = self.send([self.files])
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error_code"], "not_regular_file")
|
||||
if Path("/dev/null").exists():
|
||||
result = self.send(["/dev/null"])
|
||||
self.assertEqual(result["error_code"], "not_regular_file")
|
||||
|
||||
def test_transport_failure_is_safe_and_structured(self):
|
||||
secret_path = self.files / "private-name.txt"
|
||||
secret_path.write_bytes(b"x")
|
||||
with patch.object(mail, "_ComSession", side_effect=RuntimeError(
|
||||
"token and /private/path must not leak")):
|
||||
result = self.send([secret_path])
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error_code"], "send_unknown")
|
||||
self.assertNotIn("private", result["error"])
|
||||
self.assertNotIn("token", result["error"])
|
||||
|
||||
def test_legacy_client_send_payload_is_unchanged(self):
|
||||
app = MagicMock()
|
||||
mail._send_mail(app, SENDER, RECIPIENT, "旧主题", "旧正文")
|
||||
app.post.assert_called_once_with("/send", {
|
||||
"from_address": SENDER,
|
||||
"to_address": RECIPIENT,
|
||||
"subject": "旧主题",
|
||||
"body_text": "旧正文",
|
||||
})
|
||||
|
||||
|
||||
class BridgeAttachmentTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.flask = Flask(__name__)
|
||||
self.flask.config.update(TESTING=True)
|
||||
admin_com.install(self.flask)
|
||||
self.client = self.flask.test_client()
|
||||
self.app = MagicMock()
|
||||
self.session_patch = patch.object(backend, "_ComSession")
|
||||
session = self.session_patch.start()
|
||||
session.return_value.__enter__.return_value = self.app
|
||||
self.account_patch = patch.object(backend, "_get_account")
|
||||
self.account_patch.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.account_patch.stop()
|
||||
self.session_patch.stop()
|
||||
|
||||
def post(self, payload, endpoint=ENDPOINT):
|
||||
return self.client.post(endpoint, json=payload, headers=AUTH)
|
||||
|
||||
def test_new_endpoint_decodes_attachments_and_preserves_chinese_name(self):
|
||||
content = "中文附件".encode("utf-8")
|
||||
with patch.object(backend, "_send_mail_with_attachments", return_value=True) as send:
|
||||
response = self.post(outgoing_payload(
|
||||
attachments=[wire_attachment(content=content)]))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.get_json(), {
|
||||
"ok": True,
|
||||
"accepted": True,
|
||||
"sent_saved": True,
|
||||
"attachment_names": ["测试.txt"],
|
||||
"size": len(content),
|
||||
})
|
||||
args = send.call_args.args
|
||||
self.assertEqual(args[1:5], (SENDER, RECIPIENT, "附件测试", "正文"))
|
||||
self.assertEqual(args[5][0]["filename"], "测试.txt")
|
||||
self.assertEqual(args[5][0]["content"], content)
|
||||
|
||||
def test_new_endpoint_rejects_paths_bad_mime_and_bad_base64(self):
|
||||
bad_attachments = [
|
||||
dict(wire_attachment(), filename="../secret.txt"),
|
||||
dict(wire_attachment(), filename="C:\\secret.txt"),
|
||||
dict(wire_attachment(), content_type="text/plain\r\nX-Evil: yes"),
|
||||
dict(wire_attachment(), content_type="text/plain; charset=utf-8"),
|
||||
dict(wire_attachment(), data_base64="not base64!"),
|
||||
dict(wire_attachment(), path="C:/secret.txt"),
|
||||
]
|
||||
with patch.object(backend, "_send_mail_with_attachments") as send:
|
||||
for attachment in bad_attachments:
|
||||
with self.subTest(attachment=attachment):
|
||||
response = self.post(outgoing_payload(attachments=[attachment]))
|
||||
self.assertEqual(response.status_code, 400)
|
||||
self.assertEqual(response.get_json(), {"error": "invalid request"})
|
||||
send.assert_not_called()
|
||||
|
||||
def test_new_endpoint_rejects_empty_attachment_list_without_sending(self):
|
||||
with patch.object(backend, "_send_mail_with_attachments") as send:
|
||||
response = self.post(outgoing_payload(attachments=[]))
|
||||
self.assertEqual(response.status_code, 400)
|
||||
send.assert_not_called()
|
||||
|
||||
def test_new_endpoint_revalidates_count_file_and_total_size(self):
|
||||
with patch.object(admin_com, "MAX_ATTACHMENT_COUNT", 1), \
|
||||
patch.object(backend, "_send_mail_with_attachments") as send:
|
||||
response = self.post(outgoing_payload(attachments=[
|
||||
wire_attachment(filename="a.txt"), wire_attachment(filename="b.txt")]))
|
||||
self.assertEqual(response.status_code, 400)
|
||||
send.assert_not_called()
|
||||
|
||||
with patch.object(admin_com, "MAX_ATTACHMENT_BYTES", 4), \
|
||||
patch.object(backend, "_send_mail_with_attachments") as send:
|
||||
response = self.post(outgoing_payload(
|
||||
attachments=[wire_attachment(content=b"12345")]))
|
||||
self.assertEqual(response.status_code, 400)
|
||||
send.assert_not_called()
|
||||
|
||||
with patch.object(admin_com, "MAX_ATTACHMENT_BYTES", 4), \
|
||||
patch.object(admin_com, "MAX_ATTACHMENT_TOTAL_BYTES", 5), \
|
||||
patch.object(backend, "_send_mail_with_attachments") as send:
|
||||
response = self.post(outgoing_payload(attachments=[
|
||||
wire_attachment(content=b"123", filename="a.txt"),
|
||||
wire_attachment(content=b"456", filename="b.txt")]))
|
||||
self.assertEqual(response.status_code, 400)
|
||||
send.assert_not_called()
|
||||
|
||||
def test_request_body_limits_are_endpoint_specific(self):
|
||||
# 大于旧端点 1 MiB、小于附件端点默认 40 MiB 的合法请求可通过入口限制。
|
||||
large_content = b"x" * (800 * 1024)
|
||||
with patch.object(backend, "_send_mail_with_attachments", return_value=True):
|
||||
response = self.post(outgoing_payload(
|
||||
attachments=[wire_attachment(content=large_content)]))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
response = self.client.post(
|
||||
SEND_ENDPOINT, data=b"x" * (1024 * 1024 + 1),
|
||||
content_type="application/json", headers=AUTH)
|
||||
self.assertEqual(response.status_code, 413)
|
||||
|
||||
# 用小上限 mock 大请求，验证附件端点自身也有独立硬限制。
|
||||
with patch.object(admin_com, "ATTACHMENT_REQUEST_MAX_BYTES", 1024):
|
||||
response = self.client.post(
|
||||
ENDPOINT, data=b"x" * 1025,
|
||||
content_type="application/json", headers=AUTH)
|
||||
self.assertEqual(response.status_code, 413)
|
||||
|
||||
def test_legacy_send_endpoint_contract_is_unchanged(self):
|
||||
payload = outgoing_payload()
|
||||
payload.pop("attachments")
|
||||
with patch.object(backend, "_send_mail") as old_send, \
|
||||
patch.object(backend, "_send_mail_with_attachments") as new_send:
|
||||
response = self.post(payload, endpoint=SEND_ENDPOINT)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.get_json(), {"accepted": True})
|
||||
old_send.assert_called_once_with(
|
||||
self.app, SENDER, RECIPIENT, "附件测试", "正文")
|
||||
new_send.assert_not_called()
|
||||
|
||||
|
||||
class BackendMimeTests(unittest.TestCase):
|
||||
def test_mime_contains_utf8_body_and_rfc2231_chinese_filename(self):
|
||||
content = "二进制\x00内容".encode("utf-8")
|
||||
msg = backend._build_attachment_mime(
|
||||
SENDER, RECIPIENT, "中文主题", "纯文本正文",
|
||||
[{"filename": "中文 报告.txt", "content_type": "text/plain",
|
||||
"content": content}])
|
||||
parsed = email.message_from_bytes(msg.as_bytes())
|
||||
self.assertEqual(parsed.get_content_type(), "multipart/mixed")
|
||||
self.assertEqual(parsed["To"], RECIPIENT)
|
||||
attachment = None
|
||||
plain = None
|
||||
for part in parsed.walk():
|
||||
if part.get_content_disposition() == "attachment":
|
||||
attachment = part
|
||||
elif part.get_content_type() == "text/plain":
|
||||
plain = part
|
||||
self.assertIsNotNone(attachment)
|
||||
self.assertEqual(attachment.get_filename(), "中文 报告.txt")
|
||||
self.assertEqual(attachment.get_payload(decode=True), content)
|
||||
self.assertEqual(plain.get_payload(decode=True).decode("utf-8"), "纯文本正文")
|
||||
|
||||
def test_smtp_delivery_is_mocked_and_uses_raw_mime(self):
|
||||
smtp = MagicMock()
|
||||
smtp.sendmail.return_value = {}
|
||||
manager = MagicMock()
|
||||
manager.__enter__.return_value = smtp
|
||||
with patch.object(backend.smtplib, "SMTP", return_value=manager) as factory, \
|
||||
patch.dict(backend.CONFIG, {"smtp_host": "127.0.0.1", "smtp_port": 25,
|
||||
"smtp_ssl": False, "smtp_starttls": False,
|
||||
"smtp_username": ""}, clear=False):
|
||||
backend._deliver_mime(SENDER, RECIPIENT, b"MIME-BYTES")
|
||||
factory.assert_called_once_with("127.0.0.1", 25, timeout=30.0)
|
||||
smtp.sendmail.assert_called_once_with(SENDER, [RECIPIENT], b"MIME-BYTES")
|
||||
|
||||
def test_no_attachments_reuses_legacy_send_and_sent(self):
|
||||
app = MagicMock()
|
||||
with patch.object(backend, "_send_mail") as old_send, \
|
||||
patch.object(backend, "_save_to_sent") as old_sent, \
|
||||
patch.object(backend, "_deliver_mime") as smtp:
|
||||
result = backend._send_mail_with_attachments(
|
||||
app, SENDER, RECIPIENT, "主题", "正文", [])
|
||||
self.assertIsNone(result)
|
||||
old_send.assert_called_once_with(app, SENDER, RECIPIENT, "主题", "正文")
|
||||
old_sent.assert_called_once_with(app, SENDER, RECIPIENT, "主题", "正文")
|
||||
smtp.assert_not_called()
|
||||
|
||||
def test_sent_copy_uses_documented_com_attachment_add(self):
|
||||
message = MagicMock()
|
||||
added = []
|
||||
|
||||
def add_attachment(path):
|
||||
added.append((os.path.basename(path), Path(path).read_bytes()))
|
||||
|
||||
message.Attachments.Add.side_effect = add_attachment
|
||||
account = MagicMock()
|
||||
account.IMAPFolders.ItemByName.return_value.Messages.Add.return_value = message
|
||||
with patch.object(backend, "_get_account", return_value=account):
|
||||
saved = backend._save_attachment_copy_to_sent(
|
||||
MagicMock(), SENDER, RECIPIENT, "主题", "正文",
|
||||
[{"filename": "中文.txt", "content_type": "text/plain",
|
||||
"content": b"content"}])
|
||||
self.assertTrue(saved)
|
||||
self.assertEqual(added, [("中文.txt", b"content")])
|
||||
message.AddRecipient.assert_called_once_with("", RECIPIENT)
|
||||
message.Save.assert_called_once_with()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import mail
|
||||
|
||||
|
||||
class MailReferenceTest(unittest.TestCase):
|
||||
def test_latest_event_reference_and_explicit_id(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
history = Path(tmp) / "chat" / "history"
|
||||
history.mkdir(parents=True)
|
||||
(history / "1.txt").write_text(
|
||||
"[event:mail:2644] old\n[event:mail:2645] newest\n", encoding="utf-8")
|
||||
with patch.object(mail, "BASE_DIR", tmp):
|
||||
self.assertEqual(
|
||||
mail.resolve_context_reference(1, "帮我查看一下这篇邮件说了什么"),
|
||||
"帮我查看一下编号为 2645 的邮件说了什么")
|
||||
self.assertEqual(
|
||||
mail.resolve_context_reference(1, "查看编号2644的邮件"),
|
||||
"查看编号2644的邮件")
|
||||
|
||||
def test_missing_history_is_not_guessed(self):
|
||||
with tempfile.TemporaryDirectory() as tmp, patch.object(mail, "BASE_DIR", tmp):
|
||||
text = "帮我查看一下这篇邮件说了什么"
|
||||
self.assertEqual(mail.resolve_context_reference(1, text), text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,74 @@
|
||||
import os,sys,tempfile,json,unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch,MagicMock
|
||||
|
||||
BASE=Path(__file__).resolve().parents[1]
|
||||
sys.path[:0]=[str(BASE/'bot-utility'),str(BASE)]
|
||||
import mail
|
||||
from tg_channel import TGChannel
|
||||
|
||||
|
||||
class MigrationTests(unittest.TestCase):
|
||||
def test_mail_session_requires_tls(self):
|
||||
with patch.dict(os.environ,{'MAIL_COM_URL':'http://remote.example/admin/com'}):
|
||||
with self.assertRaises(RuntimeError):
|
||||
with mail._ComSession():pass
|
||||
|
||||
def test_mail_client_request_and_close(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p=Path(d)/'token';p.write_text('t'*128)
|
||||
with patch.dict(os.environ,{'MAIL_COM_TOKEN_FILE':str(p)}),patch('requests.Session') as factory:
|
||||
factory.return_value.post.return_value.json.return_value={'items':[{'id':42}]}
|
||||
with mail._ComSession() as session:
|
||||
result=mail._list_inbox_messages(session,'a@rrrosmontis.icu',limit=5)
|
||||
self.assertEqual(result,[{'id':42}])
|
||||
self.assertFalse(factory.return_value.trust_env)
|
||||
self.assertEqual(factory.return_value.post.call_args.kwargs['json']['limit'],5)
|
||||
factory.return_value.close.assert_called_once()
|
||||
|
||||
def test_body_attachment_metadata_does_not_open_local_eml(self):
|
||||
app=MagicMock()
|
||||
app.post.return_value={'message':{'ID':8,'Subject':'s','From':'a@b.c','Date':'d','Body':'text','HTMLBody':'<p>body</p><script>bad()</script>','image_count':2,'attachment_names':['a.pdf']}}
|
||||
m=mail._find_message_by_id(app,['a@rrrosmontis.icu'],8)
|
||||
self.assertEqual(mail._analyze_message(m),(2,['a.pdf']))
|
||||
self.assertEqual(mail._extract_body(m),'body')
|
||||
|
||||
def test_initial_poll_keeps_com_id_and_does_not_push_history(self):
|
||||
with tempfile.TemporaryDirectory() as d,patch.object(mail,'MAIL_DIR',d),patch.object(mail,'_ComSession'),patch.object(mail,'_mail_changes',return_value={'max_id':2386,'items':[]}):
|
||||
(Path(d)/'9000000001.json').write_text(json.dumps({'addresses':['a@rrrosmontis.icu'],'push':{'enabled':True,'last_seen':{}}}))
|
||||
mail._poll_next.clear()
|
||||
with patch.object(mail.notification_queue,'put') as put:
|
||||
mail._poll_once();put.assert_not_called()
|
||||
self.assertEqual(mail._load_user(9000000001)['push']['last_seen']['a@rrrosmontis.icu'],2386)
|
||||
|
||||
def test_incremental_poll_queues_without_advancing_watermark(self):
|
||||
addr='a@rrrosmontis.icu'
|
||||
with tempfile.TemporaryDirectory() as d,patch.object(mail,'MAIL_DIR',d),patch.object(mail,'_ComSession'),patch.object(mail.identity,'push_targets',return_value=[('tg',1)]):
|
||||
(Path(d)/'9000000001.json').write_text(json.dumps({'addresses':[addr],'push':{'enabled':True,'last_seen':{addr:10}}}))
|
||||
mail._poll_next.clear();mail._inflight.clear()
|
||||
changes={'max_id':99,'items':[{'id':12,'subject':'s','from':'a','from_address':'a@b.c','ts':1},{'id':14,'subject':'s','from':'a','from_address':'a@b.c','ts':1}]}
|
||||
with patch.object(mail,'_mail_changes',return_value=changes),patch.object(mail.notification_queue,'put') as put:
|
||||
mail._poll_once()
|
||||
self.assertEqual([c.args[0]['mid'] for c in put.call_args_list],[12,14])
|
||||
self.assertEqual(mail._load_user(9000000001)['push']['last_seen'][addr],10)
|
||||
mail._inflight.clear()
|
||||
|
||||
def test_oss_uris_and_traversal(self):
|
||||
self.assertEqual(TGChannel._translate_oss_path('file:///C:/oss/pixiv/a%20b.jpg'),'/mnt/oss/pixiv/a b.jpg')
|
||||
self.assertEqual(TGChannel._translate_oss_path('file:///mnt/A/pixiv/a.jpg'),'/mnt/oss/pixiv/a.jpg')
|
||||
with self.assertRaises(ValueError):TGChannel._translate_oss_path('file:///mnt/A/../secret')
|
||||
|
||||
def test_no_windows_admin_secret_or_windows_paths(self):
|
||||
self.assertIsNone(mail.HMAIL_ADMIN_PASSWORD)
|
||||
for name in ['jm','pixiv','setu']:
|
||||
module=__import__(name)
|
||||
# Runtime proxy config is private and intentionally absent in a clone.
|
||||
configured = os.environ.get('BOT_PROXY_CONFIG', str(BASE/'bot-utility'/'proxies.json'))
|
||||
self.assertEqual(Path(module.PROXY_CONFIG_PATH), Path(configured))
|
||||
self.assertNotIn('C:',module.PROXY_CONFIG_PATH)
|
||||
import pixiv,setu,for_admin_komari_agent
|
||||
self.assertEqual(pixiv.OSS_SAVE_DIR,setu.OSS_SAVE_DIR)
|
||||
self.assertEqual(for_admin_komari_agent.KOMARI_API_BASE,'https://status.rrrosmontis.icu')
|
||||
|
||||
|
||||
if __name__=='__main__':unittest.main()
|
||||
@@ -0,0 +1,28 @@
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
import identity
|
||||
from agent.modes import ModeStore
|
||||
|
||||
|
||||
class ModeConfigTest(unittest.TestCase):
|
||||
def test_admin_chat_other_command_and_saved_override(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = ModeStore(Path(tmp) / "modes.db")
|
||||
config = json.dumps({"chat_mode": {"default": "command", "admin": "chat", "group": "command"}})
|
||||
with patch("pathlib.Path.open", mock_open(read_data=config)), \
|
||||
patch.object(identity, "is_admin", side_effect=lambda uid: uid == 1):
|
||||
self.assertEqual(store.get_mode("tg", 1), "chat")
|
||||
self.assertEqual(store.get_mode("tg", 2), "command")
|
||||
self.assertEqual(store.get_mode("tg", 1, 123), "command")
|
||||
store.set_mode("tg", 2, "chat")
|
||||
self.assertEqual(store.get_mode("tg", 2), "chat")
|
||||
store.set_mode("tg", 1, "command")
|
||||
self.assertEqual(store.get_mode("tg", 1), "command")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Native deepseek-flash multimodal chat tests; no real API calls."""
|
||||
|
||||
import base64
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import agent_loop
|
||||
from agent import ToolResult
|
||||
from chat import engine
|
||||
|
||||
|
||||
class MultimodalContentTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_local_image_and_task_share_one_user_content(self):
|
||||
payload = b"\x89PNG\r\n\x1a\nsmall-test-image"
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "task.png"
|
||||
path.write_bytes(payload)
|
||||
content = await engine._build_user_content(
|
||||
"总结图片，并读出其中要做的任务",
|
||||
[{"url": str(path), "file_size": str(len(payload))}],
|
||||
)
|
||||
self.assertIsInstance(content, list)
|
||||
self.assertEqual(content[0]["type"], "image_url")
|
||||
self.assertEqual(content[1]["type"], "text")
|
||||
self.assertEqual(content[1]["text"], "总结图片，并读出其中要做的任务")
|
||||
data_url = content[0]["image_url"]["url"]
|
||||
self.assertTrue(data_url.startswith("data:image/png;base64,"))
|
||||
self.assertEqual(base64.b64decode(data_url.split(",", 1)[1]), payload)
|
||||
|
||||
async def test_text_only_content_is_unchanged(self):
|
||||
content = await engine._build_user_content("普通聊天", None)
|
||||
self.assertEqual(content, "普通聊天")
|
||||
|
||||
async def test_main_call_passes_same_multimodal_content_to_official_model(self):
|
||||
parts = [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,eA=="}},
|
||||
{"type": "text", "text": "读图"},
|
||||
]
|
||||
official = AsyncMock(return_value="ok")
|
||||
with patch.object(engine, "_modelscope_available", return_value=False), \
|
||||
patch.object(engine, "_build_user_content", new=AsyncMock(return_value=parts)), \
|
||||
patch.object(engine, "_call_deepseek_official", new=official):
|
||||
result = await engine.call_deepseek("system", "读图", image_infos=[{"url": "x"}])
|
||||
self.assertEqual(result, "ok")
|
||||
self.assertIs(official.await_args.args[1], parts)
|
||||
self.assertEqual(engine.MODEL_NAME, "deepseek-flash")
|
||||
|
||||
async def test_private_pending_flow_forwards_original_images_without_description_step(self):
|
||||
images = [{"url": "/tmp/original.png"}]
|
||||
reply = AsyncMock()
|
||||
engine._user_pending_messages[123] = ["按图执行"]
|
||||
engine._user_pending_images[123] = images
|
||||
try:
|
||||
with patch.object(engine, "_call_model_and_reply", new=reply):
|
||||
await engine._process_pending_messages(123, 0)
|
||||
finally:
|
||||
engine._user_pending_messages.pop(123, None)
|
||||
engine._user_pending_images.pop(123, None)
|
||||
engine._user_pending_tasks.pop(123, None)
|
||||
self.assertEqual(reply.await_args.args[2], "按图执行")
|
||||
self.assertIs(reply.await_args.kwargs["image_infos"], images)
|
||||
|
||||
|
||||
class FakeStore:
|
||||
def start(self, uid, conversation):
|
||||
return "task"
|
||||
def record(self, *args):
|
||||
pass
|
||||
def finish(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
class FakeExecutor:
|
||||
async def execute(self, name, args, ctx):
|
||||
return ToolResult(True, "ok", "done")
|
||||
|
||||
|
||||
class MultiRoundEngine:
|
||||
def __init__(self):
|
||||
self.image_kwargs = []
|
||||
self.reasoning_kwargs = []
|
||||
async def call_deepseek(self, *args, **kwargs):
|
||||
self.image_kwargs.append(kwargs.get("image_infos"))
|
||||
self.reasoning_kwargs.append(kwargs.get("reasoning"))
|
||||
return {"message": [{"content": "完成"}], "tools": []}
|
||||
def parse_json_response(self, raw):
|
||||
return raw
|
||||
|
||||
|
||||
class AgentMultimodalTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_tool_followup_keeps_original_images(self):
|
||||
images = [{"url": "/tmp/original.png"}]
|
||||
model = MultiRoundEngine()
|
||||
initial = {"message": [], "tools": [{"name": "mail", "args": {"action": "list"}}]}
|
||||
channel = type("Channel", (), {"name": "tg", "is_admin": lambda self, uid: True})()
|
||||
with patch.object(agent_loop, "AgentTaskStore", FakeStore), \
|
||||
patch.object(agent_loop, "ToolExecutor", FakeExecutor):
|
||||
result = await agent_loop.run_agent_loop(
|
||||
model, initial, "system", "按图执行", channel, "private", 1, 2,
|
||||
image_infos=images,
|
||||
)
|
||||
self.assertEqual(result["message"][0]["content"], "完成")
|
||||
self.assertEqual(model.image_kwargs, [images])
|
||||
self.assertEqual(model.reasoning_kwargs, [True])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Avoid formulaic user-intent paraphrase openings without removing address preferences."""
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / 'bot-utility'), str(BASE)]
|
||||
from chat import engine
|
||||
|
||||
|
||||
class NoParaphraseOpenerTests(unittest.TestCase):
|
||||
def test_global_style_covers_template_and_same_family(self):
|
||||
text = engine.load_instruction('style_instructions.txt')
|
||||
for phrase in ('主人这是在……呀', '原来主人是在', '看来主人是想',
|
||||
'主人这是……呢', '你这是/看来你是/所以你是',
|
||||
'无需每次先概括、解释或揣测', '不把猜测说成事实',
|
||||
'“主人”这个称呼保留并自然使用'):
|
||||
self.assertIn(phrase, text)
|
||||
|
||||
def test_live_and_builtin_defaults_match(self):
|
||||
text = engine.load_instruction('style_instructions.txt')
|
||||
self.assertEqual(text.strip(), engine.DEFAULT_STYLE_INSTRUCTIONS.strip())
|
||||
self.assertIn('主人这是在……呀', engine.BASE_PROMPT)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,296 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
BACKEND = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BACKEND / "bot-utility"), str(BACKEND)]
|
||||
|
||||
import notebook as notebook_module
|
||||
from notebook import HELP_TEXT, Notebook, SHANGHAI, format_memo, format_todo
|
||||
|
||||
|
||||
NOW = datetime(2030, 1, 2, 12, 0, 0, tzinfo=SHANGHAI).timestamp()
|
||||
HOUR = 3600
|
||||
DAY = 24 * HOUR
|
||||
|
||||
|
||||
def private_ctx(uid=1, channel="telegram"):
|
||||
return SimpleNamespace(user_id=uid, channel=channel, msg_type="private", target_id=uid)
|
||||
|
||||
|
||||
def group_ctx(target_id=100, channel="telegram", msg_type="group", uid=1):
|
||||
return SimpleNamespace(
|
||||
user_id=uid,
|
||||
channel=SimpleNamespace(name=channel),
|
||||
msg_type=msg_type,
|
||||
target_id=target_id,
|
||||
)
|
||||
|
||||
|
||||
def iso_at(epoch):
|
||||
return datetime.fromtimestamp(epoch, SHANGHAI).isoformat()
|
||||
|
||||
|
||||
class NotebookTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tempdir.name, "notebook-test.db")
|
||||
self.nb = Notebook(self.db_path)
|
||||
self.ctx = private_ctx(1)
|
||||
|
||||
def tearDown(self):
|
||||
self.nb.close()
|
||||
self.tempdir.cleanup()
|
||||
|
||||
def add_todo(self, content, due, expires=None, created_at=NOW, ctx=None):
|
||||
with mock.patch("notebook.time.time", return_value=created_at):
|
||||
record = self.nb.add_todo(
|
||||
ctx or self.ctx,
|
||||
content,
|
||||
iso_at(due),
|
||||
None if expires is None else iso_at(expires),
|
||||
)
|
||||
self.assertEqual(record["content"], content)
|
||||
self.assertEqual(record["status"], "active")
|
||||
return record["id"]
|
||||
|
||||
def test_memo_add_query_delete_and_content_limit(self):
|
||||
with mock.patch("notebook.time.time", side_effect=[NOW, NOW + 1]):
|
||||
first_record = self.nb.add_memo(self.ctx, "购买牛奶")
|
||||
second_record = self.nb.add_memo(self.ctx, "购买咖啡豆")
|
||||
first = first_record["id"]
|
||||
second = second_record["id"]
|
||||
self.assertEqual(first_record["content"], "购买牛奶")
|
||||
self.assertEqual(first_record["created_at"], NOW)
|
||||
self.assertRegex(first, r"^[0-9a-f]{8}$")
|
||||
self.assertNotEqual(first, second)
|
||||
self.assertEqual([m["id"] for m in self.nb.list_memos(self.ctx)], [second, first])
|
||||
self.assertEqual(
|
||||
[m["id"] for m in self.nb.list_memos(self.ctx, "牛奶")], [first]
|
||||
)
|
||||
self.assertEqual(self.nb.list_memos(self.ctx, limit=0), [])
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.add_memo(self.ctx, "x" * 2001)
|
||||
self.assertEqual(self.nb.delete_memo(self.ctx, first[:4]), first_record)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_memo(self.ctx, first)
|
||||
self.assertEqual([m["id"] for m in self.nb.list_memos(self.ctx)], [second])
|
||||
with sqlite3.connect(self.db_path) as con:
|
||||
self.assertEqual(
|
||||
con.execute("SELECT COUNT(*) FROM memo_notes WHERE id=?", (first,)).fetchone()[0],
|
||||
0,
|
||||
)
|
||||
|
||||
def test_memo_permissions_and_scope_isolation(self):
|
||||
private_id = self.nb.add_memo(private_ctx(1, channel="one"), "私聊记录")["id"]
|
||||
# 私聊只按 user_id 隔离，channel/target_id 不参与私聊作用域。
|
||||
self.assertEqual(
|
||||
self.nb.list_memos(private_ctx(1, channel="two"))[0]["id"], private_id
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_memo(private_ctx(2), private_id)
|
||||
|
||||
group_one = group_ctx(target_id=100, uid=1)
|
||||
same_group_other_sender = group_ctx(target_id=100, uid=999)
|
||||
group_two = group_ctx(target_id=200, uid=1)
|
||||
group_record = self.nb.add_memo(group_one, "群记录")
|
||||
group_id = group_record["id"]
|
||||
self.assertEqual(self.nb.list_memos(same_group_other_sender)[0]["id"], group_id)
|
||||
self.assertEqual(self.nb.list_memos(group_two), [])
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_memo(group_two, group_id)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_memo(group_ctx(channel="other"), group_id)
|
||||
self.assertEqual(
|
||||
self.nb.delete_memo(same_group_other_sender, group_id), group_record
|
||||
)
|
||||
|
||||
def test_id_prefix_must_be_unique_within_scope(self):
|
||||
with mock.patch(
|
||||
"notebook.secrets.token_hex", side_effect=["aaaaaaa1", "aaaaaaa2"]
|
||||
):
|
||||
first = self.nb.add_memo(self.ctx, "一")
|
||||
self.nb.add_memo(self.ctx, "二")
|
||||
with self.assertRaisesRegex(ValueError, "不唯一"):
|
||||
self.nb.delete_memo(self.ctx, "aaaa")
|
||||
# 其他会话中相同前缀不应影响本会话解析。
|
||||
other = private_ctx(2)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_memo(other, "aaaa")
|
||||
self.assertEqual(self.nb.delete_memo(self.ctx, first["id"]), first)
|
||||
|
||||
def test_todo_due_default_expiry_and_validation(self):
|
||||
due = NOW + 2 * HOUR
|
||||
item_id = self.add_todo("提交报告", due)
|
||||
item = self.nb.list_todos(self.ctx)[0]
|
||||
self.assertEqual(item["id"], item_id)
|
||||
self.assertEqual(item["due_at"], due)
|
||||
self.assertEqual(item["expires_at"], due + DAY)
|
||||
self.assertEqual(item["status"], "active")
|
||||
self.assertIsNone(item["completed_at"])
|
||||
|
||||
explicit_id = self.add_todo("显式过期", NOW + 3 * HOUR, NOW + 4 * HOUR)
|
||||
explicit = next(t for t in self.nb.list_todos(self.ctx) if t["id"] == explicit_id)
|
||||
self.assertEqual(explicit["expires_at"], NOW + 4 * HOUR)
|
||||
with self.assertRaisesRegex(ValueError, "未来"):
|
||||
self.add_todo("过去", NOW - 1)
|
||||
with self.assertRaisesRegex(ValueError, "晚于"):
|
||||
self.add_todo("过期等于到期", NOW + HOUR, NOW + HOUR)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.add_todo(self.ctx, "无效时间", "not-a-time")
|
||||
with self.assertRaises(ValueError):
|
||||
self.add_todo("x" * 2001, NOW + HOUR)
|
||||
|
||||
def test_complete_delete_and_permissions(self):
|
||||
item_id = self.add_todo("待办权限", NOW + HOUR)
|
||||
stranger = private_ctx(2)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.complete_todo(stranger, item_id)
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.delete_todo(stranger, item_id)
|
||||
self.assertEqual(self.nb.list_todos(self.ctx)[0]["status"], "active")
|
||||
|
||||
with mock.patch("notebook.time.time", return_value=NOW + 10):
|
||||
completed_record = self.nb.complete_todo(self.ctx, item_id[:5])
|
||||
self.assertEqual(completed_record["status"], "completed")
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.complete_todo(self.ctx, item_id)
|
||||
self.assertEqual(self.nb.list_todos(self.ctx), [])
|
||||
completed = self.nb.list_todos(self.ctx, include_completed=True)[0]
|
||||
self.assertEqual(completed["status"], "completed")
|
||||
self.assertEqual(completed["completed_at"], NOW + 10)
|
||||
self.assertEqual(self.nb.delete_todo(self.ctx, item_id)["id"], item_id)
|
||||
self.assertEqual(self.nb.list_todos(self.ctx, include_completed=True), [])
|
||||
with sqlite3.connect(self.db_path) as con:
|
||||
self.assertEqual(con.execute("SELECT COUNT(*) FROM todo_items").fetchone()[0], 0)
|
||||
|
||||
def test_cleanup_physically_deletes_only_expired_active(self):
|
||||
expired_id = self.add_todo("过期未完成", NOW + HOUR) # expires = NOW + 25h
|
||||
completed_id = self.add_todo("过期已完成", NOW + HOUR)
|
||||
active_id = self.add_todo("仍有效", NOW + 48 * HOUR)
|
||||
self.assertTrue(self.nb.complete_todo(self.ctx, completed_id))
|
||||
|
||||
# 等于 expires_at 即过期；completed 即使过期也保留。
|
||||
self.assertEqual(self.nb.cleanup_expired(NOW + 25 * HOUR), 1)
|
||||
with sqlite3.connect(self.db_path) as con:
|
||||
rows = dict(con.execute("SELECT id,status FROM todo_items").fetchall())
|
||||
self.assertNotIn(expired_id, rows)
|
||||
self.assertEqual(rows[completed_id], "completed")
|
||||
self.assertEqual(rows[active_id], "active")
|
||||
self.assertEqual(self.nb.cleanup_expired(NOW + 25 * HOUR), 0)
|
||||
|
||||
def test_expired_todo_cannot_be_completed_to_escape_deletion(self):
|
||||
item_id = self.add_todo(
|
||||
"过期后不能完成", NOW - 2 * HOUR, NOW - HOUR, created_at=NOW - 2 * DAY
|
||||
)
|
||||
with mock.patch("notebook.time.time", return_value=NOW):
|
||||
with self.assertRaises(ValueError):
|
||||
self.nb.complete_todo(self.ctx, item_id)
|
||||
with sqlite3.connect(self.db_path) as con:
|
||||
self.assertEqual(
|
||||
con.execute("SELECT COUNT(*) FROM todo_items WHERE id=?", (item_id,)).fetchone()[0],
|
||||
0,
|
||||
)
|
||||
|
||||
def test_context_limits_and_prominent_background_rule(self):
|
||||
for index in range(4):
|
||||
with mock.patch("notebook.time.time", return_value=NOW + index):
|
||||
self.nb.add_memo(self.ctx, f"备忘录{index}")
|
||||
for index in range(4):
|
||||
self.add_todo(f"待办{index}", NOW + (index + 1) * HOUR)
|
||||
|
||||
context = self.nb.context_for_chat(1, now=NOW)
|
||||
self.assertIn("⚠️ 重要规则", context)
|
||||
self.assertIn("这些仅是背景", context)
|
||||
self.assertIn("只有用户询问或直接相关时才可参考", context)
|
||||
self.assertIn("不主动重复、汇报或改变话题", context)
|
||||
self.assertIn("当前聊天优先，不能喧宾夺主", context)
|
||||
self.assertNotIn("备忘录0", context)
|
||||
for index in range(1, 4):
|
||||
self.assertIn(f"备忘录{index}", context)
|
||||
for index in range(3):
|
||||
self.assertIn(f"待办{index}", context)
|
||||
self.assertNotIn("待办3", context)
|
||||
self.assertEqual(context.count("｜到期 "), 3)
|
||||
|
||||
def test_context_overdue_72_hour_rule_query_and_physical_cleanup(self):
|
||||
created = NOW - 10 * DAY
|
||||
self.add_todo("逾期任务", NOW - HOUR, NOW + DAY, created)
|
||||
self.add_todo("边界任务", NOW + 72 * HOUR, NOW + 73 * HOUR, created)
|
||||
self.add_todo("过远任务", NOW + 73 * HOUR, NOW + 74 * HOUR, created)
|
||||
expired_id = self.add_todo("过期任务", NOW - 2 * HOUR, NOW - HOUR, created)
|
||||
completed_id = self.add_todo("已完成任务", NOW + HOUR, NOW + DAY, created)
|
||||
self.assertTrue(self.nb.complete_todo(self.ctx, completed_id))
|
||||
|
||||
context = self.nb.context_for_chat(1, now=NOW)
|
||||
self.assertIn("【已逾期】逾期任务", context)
|
||||
self.assertIn("边界任务", context)
|
||||
self.assertNotIn("过远任务", context)
|
||||
self.assertNotIn("过期任务", context)
|
||||
self.assertNotIn("已完成任务", context)
|
||||
with sqlite3.connect(self.db_path) as con:
|
||||
self.assertEqual(
|
||||
con.execute("SELECT COUNT(*) FROM todo_items WHERE id=?", (expired_id,)).fetchone()[0],
|
||||
0,
|
||||
)
|
||||
self.assertIn("逾期任务", self.nb.context_for_chat(1, "逾期", now=NOW))
|
||||
self.assertIn("逾期任务", self.nb.context_for_chat(1, "不相关的整句聊天", now=NOW))
|
||||
|
||||
def test_context_private_isolation_and_empty_result(self):
|
||||
self.nb.add_memo(self.ctx, "用户一的记录")
|
||||
self.nb.add_memo(group_ctx(target_id=100, uid=1), "群记录")
|
||||
self.assertEqual(self.nb.context_for_chat(2, now=NOW), "")
|
||||
self.assertNotIn("群记录", self.nb.context_for_chat(1, now=NOW))
|
||||
|
||||
def test_restart_persistence(self):
|
||||
memo_id = self.nb.add_memo(self.ctx, "持久备忘录")
|
||||
todo_id = self.add_todo("持久待办", NOW + DAY)
|
||||
self.nb.close()
|
||||
|
||||
restarted = Notebook(self.db_path)
|
||||
try:
|
||||
self.assertEqual(restarted.list_memos(self.ctx)[0]["id"], memo_id["id"])
|
||||
self.assertEqual(restarted.list_todos(self.ctx)[0]["id"], todo_id)
|
||||
finally:
|
||||
restarted.close()
|
||||
|
||||
def test_formatters_and_help_distinguish_reminders(self):
|
||||
with mock.patch("notebook.time.time", return_value=NOW):
|
||||
memo_id = self.nb.add_memo(self.ctx, "格式")
|
||||
memo = self.nb.list_memos(self.ctx)[0]
|
||||
self.assertIn(memo_id["id"], format_memo(memo))
|
||||
todo_id = self.add_todo("格式待办", NOW + HOUR)
|
||||
todo = self.nb.list_todos(self.ctx)[0]
|
||||
self.assertIn(todo_id, format_todo(todo, now=NOW))
|
||||
self.assertIn("2030-01-02 13:00", format_todo(todo, now=NOW))
|
||||
self.assertIn("已逾期", format_todo(todo, now=NOW + 2 * HOUR))
|
||||
self.assertIn("不会到点发送通知", HELP_TEXT)
|
||||
self.assertIn("定时提醒", HELP_TEXT)
|
||||
|
||||
|
||||
class CleanupWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_worker_catches_cleanup_exception_and_continues(self):
|
||||
nb = Notebook(":memory:")
|
||||
with (
|
||||
mock.patch.object(
|
||||
nb, "cleanup_expired", side_effect=[RuntimeError("db busy"), 0]
|
||||
) as cleanup,
|
||||
mock.patch(
|
||||
"notebook.asyncio.sleep", side_effect=[None, asyncio.CancelledError]
|
||||
),
|
||||
):
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await nb.cleanup_worker(interval=0)
|
||||
self.assertEqual(cleanup.call_count, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Notebook tool adapter and lightweight chat-context integration tests."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from agent import ToolContext, ToolExecutor
|
||||
from chat import engine
|
||||
|
||||
|
||||
class FakeNotebook:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
def add_memo(self, ctx, content):
|
||||
self.calls.append(("add_memo", content))
|
||||
return {"id": "memo0001", "content": content}
|
||||
def list_memos(self, ctx, query=None):
|
||||
return [{"id": "memo0001", "content": "钥匙在抽屉"}]
|
||||
def delete_memo(self, ctx, item_id):
|
||||
return {"id": "memo0001", "content": "钥匙在抽屉"}
|
||||
def add_todo(self, ctx, content, due_at, expires_at=None):
|
||||
self.calls.append(("add_todo", content, due_at, expires_at))
|
||||
return {"id": "todo0001", "content": content, "due_at": due_at}
|
||||
def list_todos(self, ctx, include_completed=False):
|
||||
return [{"id": "todo0001", "content": "交报告", "due_at": "tomorrow"}]
|
||||
def complete_todo(self, ctx, item_id):
|
||||
return {"id": "todo0001", "content": "交报告"}
|
||||
def delete_todo(self, ctx, item_id):
|
||||
return {"id": "todo0001", "content": "交报告"}
|
||||
|
||||
|
||||
class NotebookToolTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.fake = FakeNotebook()
|
||||
module = SimpleNamespace(
|
||||
notebook=self.fake,
|
||||
format_memo=lambda item: f"{item['id']} {item['content']}",
|
||||
format_todo=lambda item: f"{item['id']} {item['content']}",
|
||||
)
|
||||
self.executor = ToolExecutor(modules={"memo": module, "todo": module})
|
||||
self.ctx = ToolContext(
|
||||
channel=SimpleNamespace(name="tg"), msg_type="private",
|
||||
user_id=42, target_id=10001,
|
||||
)
|
||||
|
||||
async def test_schema_exposes_memo_and_todo(self):
|
||||
names = [item["name"] for item in self.executor.prompt_schema(self.ctx)]
|
||||
self.assertIn("memo", names)
|
||||
self.assertIn("todo", names)
|
||||
|
||||
async def test_memo_add_list_delete(self):
|
||||
added = await self.executor.execute("memo", {"action": "add", "params": {"content": "钥匙在抽屉"}}, self.ctx)
|
||||
self.assertTrue(added.ok)
|
||||
listed = await self.executor.execute("memo", {"action": "list", "params": {}}, self.ctx)
|
||||
self.assertIn("钥匙在抽屉", listed.safe_message)
|
||||
deleted = await self.executor.execute("memo", {"action": "delete", "params": {"id": "memo"}}, self.ctx)
|
||||
self.assertTrue(deleted.ok)
|
||||
|
||||
async def test_todo_add_list_complete_delete(self):
|
||||
added = await self.executor.execute(
|
||||
"todo", {"action": "add", "params": {"content": "交报告", "due_at": "2026-10-31T18:00:00+08:00"}}, self.ctx
|
||||
)
|
||||
self.assertTrue(added.ok)
|
||||
listed = await self.executor.execute("todo", {"action": "list", "params": {}}, self.ctx)
|
||||
self.assertIn("交报告", listed.safe_message)
|
||||
completed = await self.executor.execute("todo", {"action": "complete", "params": {"id": "todo"}}, self.ctx)
|
||||
self.assertTrue(completed.ok)
|
||||
deleted = await self.executor.execute("todo", {"action": "delete", "params": {"id": "todo"}}, self.ctx)
|
||||
self.assertTrue(deleted.ok)
|
||||
|
||||
def test_chat_prompt_contains_bounded_notebook_background(self):
|
||||
block = "【独立备忘录与近期约定】仅背景，当前聊天优先，不能喧宾夺主"
|
||||
with patch.object(engine, "user_persona_prompt", return_value="persona"), \
|
||||
patch.object(engine, "load_user_status", return_value={"pause": False, "sleep": False}), \
|
||||
patch.object(engine.memory_store, "get_relation", return_value=None), \
|
||||
patch.object(engine.memory_store, "list_entities", return_value=[]), \
|
||||
patch.object(engine.memory_store, "recall_facts", return_value=[]), \
|
||||
patch.object(engine.memory_store, "migrate_legacy"), \
|
||||
patch.object(engine, "get_history_within_budget", return_value=[]), \
|
||||
patch.object(engine, "_sticker_prompt_summary", return_value=""), \
|
||||
patch.object(engine.notebook.notebook, "context_for_chat", return_value=block) as context:
|
||||
prompt = engine.build_system_prompt(42, query="当前聊天")
|
||||
context.assert_called_once_with(42, query="当前聊天")
|
||||
self.assertIn(block, prompt)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Notification delivery must survive chat-model outages with an explicit marker."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import notification_render
|
||||
import notifications
|
||||
from chat import engine
|
||||
|
||||
|
||||
class NotificationApiFallbackTests(unittest.IsolatedAsyncioTestCase):
|
||||
def event(self):
|
||||
return notifications.Event(
|
||||
kind="komari",
|
||||
dedupe_key="komari:api-test",
|
||||
legacy_text="[Komari] 服务器「Ama-10」当前处于离线状态。",
|
||||
uid=1,
|
||||
payload={"server": "Ama-10", "status": "offline"},
|
||||
)
|
||||
|
||||
async def test_chat_api_failure_returns_raw_notice_with_marker(self):
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value=None)):
|
||||
rendered = await notification_render.render_event(self.event(), "chat", 10001)
|
||||
self.assertEqual(
|
||||
rendered,
|
||||
"[Komari] 服务器「Ama-10」当前处于离线状态。\n\n【聊天模式 API 暂时不可用】",
|
||||
)
|
||||
|
||||
async def test_command_mode_never_gets_chat_api_marker(self):
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value=None)):
|
||||
rendered = await notification_render.render_event(self.event(), "command", 10001)
|
||||
self.assertEqual(rendered, self.event().legacy_text)
|
||||
|
||||
async def test_successful_render_has_no_api_marker(self):
|
||||
with patch.object(engine, "call_deepseek", new=AsyncMock(return_value="Ama-10 掉线了。")):
|
||||
rendered = await notification_render.render_event(self.event(), "chat", 10001)
|
||||
self.assertEqual(rendered, "Ama-10 掉线了。")
|
||||
self.assertNotIn(notification_render.API_UNAVAILABLE_NOTICE, rendered)
|
||||
|
||||
def test_marker_is_not_duplicated(self):
|
||||
text = "原文\n\n" + notification_render.API_UNAVAILABLE_NOTICE
|
||||
self.assertEqual(notification_render._with_api_notice(text, "chat"), text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,62 @@
|
||||
"""MyPixiv limited placeholder responses must never be delivered as artwork."""
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import pixiv
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, body):
|
||||
self.status_code = 200
|
||||
self.headers = {"Content-Type": "image/png"}
|
||||
self.body = body
|
||||
self.closed = False
|
||||
def iter_content(self, size):
|
||||
yield self.body[:3]
|
||||
yield self.body[3:]
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class PixivRestrictedTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_placeholder_is_deleted_and_raises_restricted(self):
|
||||
body = b"my-pixiv-limited-placeholder"
|
||||
response = FakeResponse(body)
|
||||
with tempfile.TemporaryDirectory() as directory, \
|
||||
patch.object(pixiv, "PIXIV_CACHE_DIR", directory), \
|
||||
patch.object(pixiv, "RESTRICTED_IMAGE_SHA256", {hashlib.sha256(body).hexdigest()}), \
|
||||
patch("requests.get", return_value=response):
|
||||
with self.assertRaises(pixiv.PixivRestrictedError):
|
||||
pixiv.download_image(148292005)
|
||||
self.assertEqual(list(Path(directory).iterdir()), [])
|
||||
self.assertTrue(response.closed)
|
||||
|
||||
async def test_handler_returns_permission_error_without_image(self):
|
||||
sender = AsyncMock(return_value=True)
|
||||
with patch.object(pixiv, "_load_proxy_ports", return_value=[]), \
|
||||
patch.object(pixiv, "download_image", side_effect=pixiv.PixivRestrictedError("limited")), \
|
||||
patch("channel.send_message", new=sender):
|
||||
await pixiv.handle_pixiv(None, "private", 10001, 1, ".pixiv 148292005")
|
||||
self.assertEqual(sender.await_count, 2)
|
||||
for call in sender.await_args_list:
|
||||
content = call.args[3]
|
||||
self.assertIsInstance(content, str)
|
||||
self.assertIn("MyPixiv 限定", sender.await_args.args[3])
|
||||
|
||||
def test_known_placeholder_hash_is_registered(self):
|
||||
self.assertIn(
|
||||
"caa0bd60e0a0c28a4bd80b6b7024fb4956ffff7c536f4fe9d2d2588b660524c7",
|
||||
pixiv.RESTRICTED_IMAGE_SHA256,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Prompt-cache ordering tests: stable instructions first, time in user messages."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from agent import ToolContext
|
||||
from chat import engine
|
||||
|
||||
|
||||
class PromptCacheOrderTests(unittest.IsolatedAsyncioTestCase):
|
||||
def build_private(self):
|
||||
channel = SimpleNamespace(name="tg", is_admin=lambda uid: True)
|
||||
ctx = ToolContext(channel=channel, msg_type="private", user_id=42, target_id=10001, is_admin=True)
|
||||
with patch.object(engine, "user_persona_prompt", return_value="PERSONA"), \
|
||||
patch.object(engine, "load_user_status", return_value={"pause": False, "sleep": False}), \
|
||||
patch.object(engine.memory_store, "get_relation", return_value=None), \
|
||||
patch.object(engine.memory_store, "list_entities", return_value=[]), \
|
||||
patch.object(engine.memory_store, "recall_facts", return_value=[(1, "长期事实", 5)]), \
|
||||
patch.object(engine.memory_store, "migrate_legacy"), \
|
||||
patch.object(engine, "get_history_within_budget", return_value=["最近历史"]), \
|
||||
patch.object(engine.notebook.notebook, "context_for_chat", return_value="笔记本背景"), \
|
||||
patch.object(engine, "_sticker_prompt_summary", return_value="贴纸场景"):
|
||||
return engine.build_system_prompt(42, query="当前问题", tool_context=ctx)
|
||||
|
||||
def test_prompt_budget_is_128k(self):
|
||||
self.assertEqual(engine.PROMPT_TOKEN_BUDGET, 131072)
|
||||
|
||||
def test_private_stable_instructions_precede_dynamic_sections(self):
|
||||
prompt = self.build_private()
|
||||
self.assertNotIn("现在是 ", prompt)
|
||||
agent = prompt.index("当你回复时")
|
||||
memory = prompt.index("【长期记忆】")
|
||||
status = prompt.index("当前用户状态")
|
||||
history = prompt.index("以下是最近你和用户的聊天记录")
|
||||
self.assertLess(agent, memory)
|
||||
self.assertLess(memory, status)
|
||||
self.assertLess(status, history)
|
||||
self.assertIn("你还可以调用以下 utility 工具", prompt[:history])
|
||||
|
||||
def test_sticker_prompt_prefers_meaning_over_description(self):
|
||||
expected = "理解其情绪、含义和使用场景即可，不要主动描述画面内容"
|
||||
self.assertIn(expected, engine.load_instruction("agent_instructions.txt"))
|
||||
self.assertIn(expected, engine.load_instruction("group_instructions.txt"))
|
||||
|
||||
def test_time_is_prepended_to_user_message(self):
|
||||
message = engine._message_with_time("用户输入")
|
||||
self.assertTrue(message.startswith("现在是 "))
|
||||
self.assertTrue(message.endswith("\n用户输入"))
|
||||
|
||||
async def test_group_prompt_has_stable_prefix_and_time_in_user_message(self):
|
||||
response = '{"message":[{"content":"ok"}],"tools":[],"need_reasoning":false}'
|
||||
model = AsyncMock(return_value=response)
|
||||
with patch.object(engine, "get_recent_group_history", return_value=["群历史"]), \
|
||||
patch.object(engine, "call_deepseek", new=model), \
|
||||
patch.object(engine, "send_group_message", new=AsyncMock(return_value=True)), \
|
||||
patch.object(engine, "append_group_history"):
|
||||
await engine._call_group_model_and_reply_impl(None, 1, "群消息", 0)
|
||||
system, user = model.await_args.args[:2]
|
||||
self.assertNotIn("现在是 ", system)
|
||||
self.assertTrue(user.startswith("现在是 "))
|
||||
self.assertLess(system.index("当你回复时"), system.index("以下是本群的最近聊天记录"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,85 @@
|
||||
"""Hybrid reasoning policy: casual chat off, Agent planning on."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from chat import engine
|
||||
|
||||
|
||||
class ReasoningPolicyTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_reasoning_parameters_are_explicit(self):
|
||||
self.assertEqual(
|
||||
engine._reasoning_params(False),
|
||||
{"thinking": {"type": "disabled"}},
|
||||
)
|
||||
self.assertEqual(
|
||||
engine._reasoning_params(True),
|
||||
{"thinking": {"type": "enabled"}, "reasoning_effort": "high"},
|
||||
)
|
||||
|
||||
def test_any_tool_call_requires_agent_reasoning(self):
|
||||
context = object()
|
||||
for tool in ("setu", "web_search", "fetch", "mail"):
|
||||
parsed = {"understood": True, "tools": [{"name": tool}]}
|
||||
self.assertTrue(engine._tool_plan_requires_reasoning(parsed, context))
|
||||
self.assertFalse(engine._tool_plan_requires_reasoning({"tools": []}, context))
|
||||
self.assertFalse(engine._tool_plan_requires_reasoning({"understood": False, "tools": [{}]}, context))
|
||||
self.assertFalse(engine._tool_plan_requires_reasoning({"tools": [{}]}, None))
|
||||
|
||||
def test_model_self_assessment_can_request_reasoning(self):
|
||||
context = object()
|
||||
parsed = {"need_reasoning": True, "tools": []}
|
||||
self.assertTrue(engine._chat_requires_reasoning(parsed, context, False))
|
||||
self.assertFalse(engine._chat_requires_reasoning(parsed, context, True))
|
||||
self.assertFalse(engine._chat_requires_reasoning({"need_reasoning": False, "tools": []}, context, False))
|
||||
prompt = engine.load_instruction("agent_instructions.txt")
|
||||
self.assertIn("need_reasoning", prompt)
|
||||
self.assertIn("没有把握", prompt)
|
||||
|
||||
def test_last_json_retry_uses_reasoning(self):
|
||||
self.assertFalse(engine._json_retry_uses_reasoning(0))
|
||||
self.assertTrue(engine._json_retry_uses_reasoning(engine.MAX_RETRIES - 2))
|
||||
self.assertFalse(engine._json_retry_uses_reasoning(engine.MAX_RETRIES - 1))
|
||||
|
||||
def test_reminder_tone_forbids_threatening_or_startling_users(self):
|
||||
prompt = engine.BASE_PROMPT
|
||||
self.assertIn("威胁、恐吓、惊吓", prompt)
|
||||
self.assertIn("底色必须温和、让人安心", prompt)
|
||||
self.assertIn("别被白花吓一跳", prompt)
|
||||
self.assertIn("只响这一次", prompt)
|
||||
self.assertIn("只响一次", prompt)
|
||||
self.assertIn("今晚八点提醒你吃饭", prompt)
|
||||
|
||||
async def test_plain_chat_defaults_to_non_thinking(self):
|
||||
official = AsyncMock(return_value="ok")
|
||||
with patch.object(engine, "_modelscope_available", return_value=False), \
|
||||
patch.object(engine, "_build_user_content", new=AsyncMock(return_value="聊天")), \
|
||||
patch.object(engine, "_call_deepseek_official", new=official):
|
||||
await engine.call_deepseek("system", "聊天")
|
||||
self.assertIs(official.await_args.kwargs["reasoning"], False)
|
||||
|
||||
async def test_agent_request_explicitly_uses_high_reasoning(self):
|
||||
official = AsyncMock(return_value="ok")
|
||||
with patch.object(engine, "_modelscope_available", return_value=False), \
|
||||
patch.object(engine, "_build_user_content", new=AsyncMock(return_value="调用工具")), \
|
||||
patch.object(engine, "_call_deepseek_official", new=official):
|
||||
await engine.call_deepseek("system", "调用工具", reasoning=True)
|
||||
self.assertIs(official.await_args.kwargs["reasoning"], True)
|
||||
|
||||
async def test_modelscope_receives_same_reasoning_flag(self):
|
||||
modelscope = AsyncMock(return_value="ok")
|
||||
with patch.object(engine, "_modelscope_available", return_value=True), \
|
||||
patch.object(engine, "_build_user_content", new=AsyncMock(return_value="聊天")), \
|
||||
patch.object(engine, "_call_modelscope", new=modelscope), \
|
||||
patch.object(engine, "_modelscope_mark_success"):
|
||||
await engine.call_deepseek("system", "聊天", reasoning=True)
|
||||
self.assertIs(modelscope.await_args.kwargs["reasoning"], True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Agent schedule-tool adapter tests with an isolated fake scheduler."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from agent import ToolContext, ToolExecutor
|
||||
|
||||
|
||||
class FakeScheduler:
|
||||
def __init__(self):
|
||||
self.created = None
|
||||
self.cancelled = None
|
||||
def create_from_params(self, ctx, params):
|
||||
self.created = (ctx, params)
|
||||
return {"id": "abc12345", "message": params["message"], "kind": params["kind"], "next_run_text": "今晚八点"}
|
||||
def list_for_context(self, ctx):
|
||||
return [{"id": "abc12345", "message": "吃饭", "kind": "once", "next_run_text": "今晚八点"}]
|
||||
def cancel_for_context(self, ctx, task_id):
|
||||
self.cancelled = (ctx, task_id)
|
||||
return {"id": "abc12345", "message": "吃饭"}
|
||||
|
||||
|
||||
class ScheduleToolTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.fake = FakeScheduler()
|
||||
module = SimpleNamespace(scheduler=self.fake, format_task=lambda task: f"{task['id']} {task['message']} {task.get('next_run_text', '')}")
|
||||
self.executor = ToolExecutor(modules={"schedule": module})
|
||||
self.ctx = ToolContext(
|
||||
channel=SimpleNamespace(name="tg"), msg_type="private",
|
||||
user_id=1, target_id=10001, is_admin=False,
|
||||
)
|
||||
|
||||
async def test_schema_exposes_schedule_to_normal_users(self):
|
||||
names = [item["name"] for item in self.executor.prompt_schema(self.ctx)]
|
||||
self.assertIn("schedule", names)
|
||||
|
||||
async def test_create_uses_current_conversation(self):
|
||||
result = await self.executor.execute(
|
||||
"schedule",
|
||||
{"action": "create", "params": {"message": "吃饭", "kind": "once", "run_at": "2026-10-30T20:00:00+08:00"}},
|
||||
self.ctx,
|
||||
)
|
||||
self.assertTrue(result.ok)
|
||||
self.assertIs(self.fake.created[0], self.ctx)
|
||||
self.assertEqual(self.fake.created[1]["message"], "吃饭")
|
||||
|
||||
async def test_list_and_cancel(self):
|
||||
listed = await self.executor.execute("schedule", {"action": "list", "params": {}}, self.ctx)
|
||||
self.assertIn("abc12345", listed.safe_message)
|
||||
cancelled = await self.executor.execute(
|
||||
"schedule", {"action": "cancel", "params": {"id": "abc"}}, self.ctx
|
||||
)
|
||||
self.assertTrue(cancelled.ok)
|
||||
self.assertEqual(self.fake.cancelled, (self.ctx, "abc"))
|
||||
|
||||
async def test_cancel_rejects_extra_params(self):
|
||||
result = await self.executor.execute(
|
||||
"schedule", {"action": "cancel", "params": {"id": "abc", "uid": 2}}, self.ctx
|
||||
)
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.code, "invalid_args")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Adversarial tests for command authorization and Telegram sticker delivery."""
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
|
||||
def load_dispatch():
|
||||
"""Import dispatch without optional network dependencies of unused modules."""
|
||||
stubs = {}
|
||||
for name in (
|
||||
"setu", "jm", "pixiv", "bilibili", "mail", "live_monitor", "ginkgostu",
|
||||
"identity", "identity_cmd", "for_admin_komari_agent", "for_admin_catch_log",
|
||||
"mode_commands",
|
||||
):
|
||||
stubs[name] = types.ModuleType(name)
|
||||
help_ui = types.ModuleType("help_ui")
|
||||
help_ui.command_index = lambda: "help"
|
||||
stubs["help_ui"] = help_ui
|
||||
deepseek = types.ModuleType("deepseek")
|
||||
deepseek.__path__ = []
|
||||
deepseek.deepseek_balance = types.ModuleType("deepseek.deepseek_balance")
|
||||
stubs["deepseek"] = deepseek
|
||||
stubs["deepseek.deepseek_balance"] = deepseek.deepseek_balance
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"dispatch_security_test", BASE / "bot-utility" / "dispatch.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
with mock.patch.dict(sys.modules, stubs):
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class ChatAdminTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_chat_on_off_requires_admin(self):
|
||||
dispatch = load_dispatch()
|
||||
channel = types.SimpleNamespace(is_admin=mock.Mock(return_value=False))
|
||||
engine = types.SimpleNamespace(set_chat_enabled=mock.Mock())
|
||||
chat = types.ModuleType("chat")
|
||||
chat.engine = engine
|
||||
dispatch.send_message = mock.AsyncMock()
|
||||
with mock.patch.dict(sys.modules, {"chat": chat}):
|
||||
handled = await dispatch._handle_chat_legacy(channel, "group", 20, 10, ".chat on")
|
||||
self.assertTrue(handled)
|
||||
engine.set_chat_enabled.assert_not_called()
|
||||
dispatch.send_message.assert_awaited_once_with(channel, "group", 20, "该指令仅管理员可用")
|
||||
|
||||
async def test_admin_can_toggle_chat(self):
|
||||
dispatch = load_dispatch()
|
||||
channel = types.SimpleNamespace(is_admin=mock.Mock(return_value=True))
|
||||
engine = types.SimpleNamespace(set_chat_enabled=mock.Mock())
|
||||
chat = types.ModuleType("chat")
|
||||
chat.engine = engine
|
||||
dispatch.send_message = mock.AsyncMock()
|
||||
with mock.patch.dict(sys.modules, {"chat": chat}):
|
||||
for arg, enabled in (("on", True), ("off", False)):
|
||||
with self.subTest(arg=arg):
|
||||
handled = await dispatch._handle_chat_legacy(
|
||||
channel, "private", 10, 10, ".chat " + arg,
|
||||
)
|
||||
self.assertTrue(handled)
|
||||
engine.set_chat_enabled.assert_called_with(enabled)
|
||||
self.assertEqual(engine.set_chat_enabled.call_count, 2)
|
||||
|
||||
|
||||
try:
|
||||
from aiogram.types import FSInputFile
|
||||
except ModuleNotFoundError:
|
||||
aiogram = types.ModuleType("aiogram")
|
||||
aiogram.__path__ = []
|
||||
aiogram_types = types.ModuleType("aiogram.types")
|
||||
|
||||
class FSInputFile:
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
|
||||
aiogram_types.FSInputFile = FSInputFile
|
||||
aiogram.types = aiogram_types
|
||||
sys.modules["aiogram"] = aiogram
|
||||
sys.modules["aiogram.types"] = aiogram_types
|
||||
|
||||
from tg_channel import TGChannel
|
||||
|
||||
|
||||
class TelegramStickerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_agent_sticker_uses_sticker_api_not_photo(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image = Path(directory) / "sticker.webp"
|
||||
image.write_bytes(b"sticker")
|
||||
bot = types.SimpleNamespace(
|
||||
send_sticker=mock.AsyncMock(return_value=object()),
|
||||
send_photo=mock.AsyncMock(return_value=object()),
|
||||
send_document=mock.AsyncMock(return_value=object()),
|
||||
)
|
||||
channel = TGChannel(bot, admin_ids=[])
|
||||
sent = await channel.send("private", 123, [{
|
||||
"type": "image", "data": {"file": image.as_uri(), "sub_type": 1},
|
||||
}])
|
||||
self.assertTrue(sent)
|
||||
bot.send_sticker.assert_awaited_once()
|
||||
self.assertEqual(bot.send_sticker.await_args.kwargs["chat_id"], 123)
|
||||
bot.send_photo.assert_not_awaited()
|
||||
bot.send_document.assert_not_awaited()
|
||||
|
||||
async def test_regular_image_still_uses_photo_api(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image = Path(directory) / "image.jpg"
|
||||
image.write_bytes(b"image")
|
||||
bot = types.SimpleNamespace(
|
||||
send_sticker=mock.AsyncMock(return_value=object()),
|
||||
send_photo=mock.AsyncMock(return_value=object()),
|
||||
send_document=mock.AsyncMock(return_value=object()),
|
||||
)
|
||||
channel = TGChannel(bot, admin_ids=[])
|
||||
sent = await channel.send("private", 123, [{
|
||||
"type": "image", "data": {"file": image.as_uri()},
|
||||
}])
|
||||
self.assertTrue(sent)
|
||||
bot.send_photo.assert_awaited_once()
|
||||
bot.send_sticker.assert_not_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,139 @@
|
||||
"""Skill orchestration tests with fake models, isolated skill directories and no sends."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / 'bot-utility'), str(BASE)]
|
||||
import skill_registry
|
||||
import skill_runtime
|
||||
import help_ui
|
||||
import notification_render
|
||||
import notifications
|
||||
import agent_loop
|
||||
from chat import engine
|
||||
|
||||
|
||||
class SkillIntegrationTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
base = Path(self.tmp.name)
|
||||
root = base / 'skills'
|
||||
self.body = 'SPECIFIC_STYLE_RULE: 用户称呼优先于默认规范。'
|
||||
for name in ('language', 'result-style'):
|
||||
folder = root / name
|
||||
folder.mkdir(parents=True)
|
||||
(folder / 'SKILL.md').write_text(
|
||||
f'---\nname: {name}\ndescription: Language style\n---\n' + self.body + name,
|
||||
encoding='utf-8')
|
||||
config = base / 'skills_config.yml'
|
||||
config.write_text(json.dumps({
|
||||
'enabled': True, 'allowlist': ['language', 'result-style'],
|
||||
'always': ['language'], 'scenes': {'tool_result': ['result-style']},
|
||||
}), encoding='utf-8')
|
||||
self.registry = skill_registry.SkillRegistry(root, config)
|
||||
self.registry_patch = patch.object(skill_runtime, 'registry', self.registry)
|
||||
self.registry_patch.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.registry_patch.stop()
|
||||
self.tmp.cleanup()
|
||||
|
||||
async def test_help_keeps_personal_settings_and_unwraps_text(self):
|
||||
model = AsyncMock(return_value=json.dumps({'text': '主人，可以查邮件。', 'tools': []}))
|
||||
with patch.object(engine, '_call_deepseek_transport', new=model), \
|
||||
patch.object(engine, 'user_persona_prompt', return_value='USER_SETTING: 称呼主人'):
|
||||
result = await help_ui.chat_help(42)
|
||||
self.assertEqual(result, '主人，可以查邮件。')
|
||||
prompt = model.await_args.args[0]
|
||||
self.assertIn(self.body, prompt)
|
||||
self.assertIn('USER_SETTING: 称呼主人', prompt)
|
||||
self.assertIn('help', prompt)
|
||||
self.assertIsNone(skill_runtime.current_turn())
|
||||
|
||||
async def test_notification_and_reminder_open_own_scene(self):
|
||||
scenes = []
|
||||
async def model(*args, **kwargs):
|
||||
scenes.append(skill_runtime.current_turn().scene)
|
||||
return json.dumps({'text': '提醒内容', 'tools': []})
|
||||
with patch.object(engine, '_call_deepseek_transport', new=model), \
|
||||
patch.object(engine, 'user_persona_prompt', return_value='USER_SETTING'):
|
||||
for kind in ('komari', 'reminder'):
|
||||
event = notifications.Event(kind=kind, dedupe_key='test-' + kind,
|
||||
uid=42, legacy_text='原文')
|
||||
self.assertEqual(await notification_render.render_event(event, 'chat', 1), '提醒内容')
|
||||
self.assertEqual(scenes, ['notification', 'reminder'])
|
||||
|
||||
async def test_notification_cannot_execute_business_tool_and_keeps_fallback(self):
|
||||
response = json.dumps({'text': '不应发送', 'tools': [{'name': 'akaros_mail', 'args': {}}]})
|
||||
with patch.object(engine, '_call_deepseek_transport', new=AsyncMock(return_value=response)), \
|
||||
patch.object(engine, 'user_persona_prompt', return_value='base'):
|
||||
event = notifications.Event(kind='komari', dedupe_key='readonly', uid=42, legacy_text='原始告警')
|
||||
result = await notification_render.render_event(event, 'chat', 1)
|
||||
self.assertIn('原始告警', result)
|
||||
self.assertNotIn('不应发送', result)
|
||||
|
||||
async def test_private_and_proactive_sessions_cover_entire_generation(self):
|
||||
seen = []
|
||||
async def impl(*args, **kwargs):
|
||||
seen.append(skill_runtime.current_turn().scene)
|
||||
with patch.object(engine, '_resolve_delivery', return_value=None), \
|
||||
patch.object(engine, '_call_model_and_reply_impl', new=impl), \
|
||||
patch.object(notifications.center, 'begin_conversation'), \
|
||||
patch.object(notifications.center, 'end_conversation'), \
|
||||
patch.object(notifications.center, 'deliver_due', new=AsyncMock()):
|
||||
await engine._call_model_and_reply(None, 101, 'hi', False)
|
||||
await engine._call_model_and_reply(None, 101, 'hi', True)
|
||||
self.assertEqual(seen, ['private_chat', 'proactive'])
|
||||
self.assertIsNone(skill_runtime.current_turn())
|
||||
|
||||
async def test_group_has_group_session_not_private_persona(self):
|
||||
seen = []
|
||||
async def impl(*args, **kwargs):
|
||||
seen.append(skill_runtime.current_turn().scene)
|
||||
with patch.object(engine, '_call_group_model_and_reply_impl', new=impl):
|
||||
await engine._call_group_model_and_reply(None, 201, 'hi', 42)
|
||||
self.assertEqual(seen, ['group_chat'])
|
||||
|
||||
async def test_agent_followup_binds_result_style_without_repeating_business(self):
|
||||
final = json.dumps({'message': [{'content': 'done'}], 'tools': []})
|
||||
model = AsyncMock(return_value=final)
|
||||
channel = SimpleNamespace(name='tg', is_admin=lambda uid: False)
|
||||
store = MagicMock()
|
||||
with skill_runtime.skill_session('private_chat'), \
|
||||
patch.object(engine, '_call_deepseek_transport', new=model), \
|
||||
patch.object(engine, 'run_tools', new=AsyncMock(return_value='read result')) as tools, \
|
||||
patch.object(agent_loop, 'AgentTaskStore', return_value=store), \
|
||||
patch.object(agent_loop, 'ToolExecutor'):
|
||||
result = await agent_loop.run_agent_loop(engine,
|
||||
{'message': [], 'tools': [{'name': 'web_search', 'args': {'query': 'x'}}]},
|
||||
'system', 'user', channel, 'private', 42, 1)
|
||||
self.assertEqual(result['tools'], [])
|
||||
tools.assert_awaited_once()
|
||||
self.assertIn(self.body + 'result-style', model.await_args.args[0])
|
||||
|
||||
async def test_internal_sticker_audit_suspends_skills(self):
|
||||
image = Path(self.tmp.name) / 'preview.png'
|
||||
image.write_bytes(b'preview')
|
||||
model = AsyncMock(return_value='{"send":true}')
|
||||
with skill_runtime.skill_session('private_chat'), \
|
||||
patch.object(engine, '_call_deepseek_transport', new=model):
|
||||
self.assertTrue(await engine._confirm_sticker({'preview_path': str(image)}, 'happy', 'hi'))
|
||||
self.assertNotIn(self.body, model.await_args.args[0])
|
||||
|
||||
def test_usage_logger_contains_counts_not_payload(self):
|
||||
payload = {'usage': {'prompt_tokens': 100, 'prompt_tokens_details': {'cached_tokens': 80}},
|
||||
'choices': [{'message': {'content': 'PRIVATE_MESSAGE'}}]}
|
||||
with patch('builtins.print') as output:
|
||||
engine._log_model_usage('test', payload)
|
||||
text = output.call_args.args[0]
|
||||
self.assertIn('cached_tokens', text)
|
||||
self.assertNotIn('PRIVATE_MESSAGE', text)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,283 @@
|
||||
"""Standalone tests: imports only the registry, never production bot modules."""
|
||||
import dataclasses
|
||||
import importlib.util
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import yaml
|
||||
|
||||
_MODULE_PATH = Path(__file__).resolve().parents[1] / "bot-utility" / "skill_registry.py"
|
||||
_SPEC = importlib.util.spec_from_file_location("_isolated_skill_registry", _MODULE_PATH)
|
||||
skills = importlib.util.module_from_spec(_SPEC)
|
||||
sys.modules[_SPEC.name] = skills
|
||||
_SPEC.loader.exec_module(skills)
|
||||
|
||||
|
||||
class SkillRegistryTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.base = Path(self.tmp.name)
|
||||
self.root = self.base / "skills"
|
||||
self.root.mkdir()
|
||||
self.config = self.base / "skills_config.yml"
|
||||
self.registry = skills.SkillRegistry(self.root, self.config)
|
||||
|
||||
def configure(self, **kwargs):
|
||||
config = {"enabled": True, "allowlist": [], "always": [], "scenes": {}}
|
||||
config.update(kwargs)
|
||||
self.config.write_text(yaml.safe_dump(config), encoding="utf-8")
|
||||
|
||||
def skill(self, name="alpha", body="Pinned instructions.\n", references=None, extra=""):
|
||||
directory = self.root / name
|
||||
directory.mkdir(exist_ok=True)
|
||||
(directory / "SKILL.md").write_text(
|
||||
f"---\nname: {name}\ndescription: Description for {name}\n{extra}---\n{body}",
|
||||
encoding="utf-8")
|
||||
for path, content in (references or {}).items():
|
||||
target = directory / "references" / path
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(content, encoding="utf-8")
|
||||
return directory
|
||||
|
||||
def error(self, code, fn, *args):
|
||||
with self.assertRaises(skills.SkillError) as caught:
|
||||
fn(*args)
|
||||
self.assertEqual(caught.exception.code, code)
|
||||
self.assertNotIn(str(self.base), str(caught.exception))
|
||||
|
||||
def test_missing_empty_and_default_config(self):
|
||||
for text in (None, "", "# empty\n", "{}\n"):
|
||||
if text is not None:
|
||||
self.config.write_text(text, encoding="utf-8")
|
||||
snapshot = self.registry.snapshot()
|
||||
self.assertTrue(snapshot.enabled)
|
||||
self.assertEqual(dict(snapshot.records), {})
|
||||
self.assertEqual(snapshot.catalog_text(), "")
|
||||
self.assertEqual(snapshot.bound_names("private_chat"), ())
|
||||
self.assertEqual(snapshot.limits["max_tokens"], 12000)
|
||||
self.assertEqual(skills.registry.root, _MODULE_PATH.parent / "skills")
|
||||
self.assertEqual(skills.registry.config_path, _MODULE_PATH.parent / "skills_config.yml")
|
||||
|
||||
def test_catalog_allowlist_bindings_and_metadata(self):
|
||||
self.skill("alpha", references={"nested/info.md": "Secret ref"},
|
||||
extra="metadata:\n author: Tester\nallowed-tools: Bash Read\n")
|
||||
self.skill("beta")
|
||||
self.skill("excluded")
|
||||
self.configure(allowlist=["beta", "alpha"], always=["alpha"],
|
||||
scenes={"notification": ["beta", "alpha"], "reminder": ["alpha"],
|
||||
"private_chat": ["beta"]})
|
||||
snapshot = self.registry.snapshot()
|
||||
self.assertEqual(tuple(snapshot.records), ("alpha", "beta"))
|
||||
self.assertEqual(snapshot.catalog_text(), "alpha: Description for alpha\nbeta: Description for beta")
|
||||
self.assertNotIn("Secret ref", snapshot.catalog_text())
|
||||
self.assertEqual(snapshot.bound_names("private_chat"), ("alpha", "beta"))
|
||||
self.assertEqual(snapshot.bound_names("reminder"), ("alpha", "beta"))
|
||||
self.assertEqual(snapshot.scenes["reminder"], ("beta", "alpha"))
|
||||
self.assertEqual(snapshot.records["alpha"].allowed_tools, ("Bash", "Read"))
|
||||
self.assertEqual(snapshot.records["alpha"].metadata["author"], "Tester")
|
||||
self.assertEqual(snapshot.read("alpha", "references/nested/info.md"), "Secret ref")
|
||||
self.error("unknown_skill", snapshot.load, "excluded")
|
||||
|
||||
def test_snapshot_pinning_immutability_and_hash(self):
|
||||
directory = self.skill(references={"info.txt": "old"})
|
||||
self.configure(allowlist=["alpha"])
|
||||
first = self.registry.snapshot()
|
||||
with self.assertRaises(TypeError):
|
||||
first.records["new"] = first.records["alpha"]
|
||||
with self.assertRaises(TypeError):
|
||||
first.records["alpha"].references["references/info.txt"] = "new"
|
||||
with self.assertRaises(TypeError):
|
||||
first.limits["max_tokens"] = 3
|
||||
with self.assertRaises(dataclasses.FrozenInstanceError):
|
||||
first.records["alpha"].body = "new"
|
||||
os.utime(directory / "SKILL.md", (1, 1))
|
||||
self.assertIs(first, self.registry.snapshot())
|
||||
other = skills.SkillRegistry(self.root, self.config).snapshot()
|
||||
self.assertEqual(first.fingerprint, other.fingerprint)
|
||||
self.skill(body="New body", references={"info.txt": "new"})
|
||||
second = self.registry.snapshot()
|
||||
self.assertNotEqual(first.fingerprint, second.fingerprint)
|
||||
self.assertEqual(first.load("alpha"), "Pinned instructions.\n")
|
||||
self.assertEqual(first.read("alpha", "references/info.txt"), "old")
|
||||
self.assertEqual(second.load("alpha"), "New body")
|
||||
self.assertEqual(second.read("alpha", "references/info.txt"), "new")
|
||||
self.configure(allowlist=["alpha"], limits={"max_tokens": 100})
|
||||
self.assertNotEqual(second.fingerprint, self.registry.snapshot().fingerprint)
|
||||
|
||||
def test_invalid_config_keeps_last_good_and_disable_wins(self):
|
||||
self.skill()
|
||||
self.configure(allowlist=["alpha"], always=["alpha"])
|
||||
good = self.registry.snapshot()
|
||||
for text in ("allowlist: [", "enabled: true\nenabled: false\n", "[]", "enabled: yesplease",
|
||||
"allowlist: [alpha, alpha]", "allowlist: [../alpha]", "always: [alpha]",
|
||||
"scenes: {unknown: []}", "limits: {max_rounds: true}",
|
||||
"allowlist: &a [alpha]\nalways: *a", "!!python/object:os.system {}",
|
||||
"unexpected: value"):
|
||||
with self.subTest(text=text):
|
||||
self.config.write_text(text, encoding="utf-8")
|
||||
self.assertIs(self.registry.snapshot(), good)
|
||||
self.config.write_text("enabled: false\nalways: [missing]\nlimits: broken\n", encoding="utf-8")
|
||||
disabled = self.registry.snapshot()
|
||||
self.assertFalse(disabled.enabled)
|
||||
self.assertEqual(disabled.catalog_text(), "")
|
||||
self.assertEqual(disabled.bound_names("help"), ())
|
||||
self.assertFalse(disabled.records)
|
||||
self.error("disabled", disabled.load, "alpha")
|
||||
self.error("disabled", disabled.read, "alpha", "references/info.md")
|
||||
self.configure(allowlist=["alpha"])
|
||||
self.assertTrue(self.registry.snapshot().enabled)
|
||||
|
||||
def test_invalid_first_config_is_fail_closed(self):
|
||||
self.config.write_text("allowlist: [", encoding="utf-8")
|
||||
self.assertFalse(self.registry.snapshot().enabled)
|
||||
|
||||
def test_limit_clamping_and_types(self):
|
||||
self.configure(limits={"body_bytes": -4, "reference_bytes": 10**9,
|
||||
"max_skills": 0, "max_rounds": 10**9, "max_tokens": -10})
|
||||
result = self.registry.snapshot()
|
||||
self.assertEqual(dict(result.limits), {"body_bytes": 1, "reference_bytes": 65536,
|
||||
"max_skills": 1, "max_rounds": 8, "max_tokens": 1})
|
||||
for limits in ({"max_tokens": "100"}, {"unknown": 1}, [], None):
|
||||
self.configure(limits=limits)
|
||||
self.assertIs(self.registry.snapshot(), result)
|
||||
|
||||
def test_invalid_skills_quarantined_independently(self):
|
||||
self.skill("valid")
|
||||
bad = self.skill("invalid")
|
||||
self.configure(allowlist=["valid", "invalid"])
|
||||
for document in ("no frontmatter", "---\nname: invalid\ndescription: [a]\n---\nbody",
|
||||
"---\nname: valid\ndescription: duplicate name\n---\nbody",
|
||||
"---\nname: invalid\nname: invalid\ndescription: desc\n---\nbody",
|
||||
"---\nname: invalid\ndescription: ''\n---\nbody",
|
||||
"---\nname: invalid\ndescription: " + "x" * 1025 + "\n---\nbody",
|
||||
"---\nname: invalid\ndescription: hi\nmetadata: {x: []}\n---\nbody"):
|
||||
with self.subTest(document=document):
|
||||
(bad / "SKILL.md").write_text(document, encoding="utf-8")
|
||||
with self.assertLogs(skills.__name__, level="WARNING") as logs:
|
||||
snapshot = self.registry.snapshot()
|
||||
self.assertEqual(tuple(snapshot.records), ("valid",))
|
||||
self.assertTrue(any("quarantined" in line for line in logs.output))
|
||||
self.assertTrue(all(str(self.base) not in line for line in logs.output))
|
||||
|
||||
def test_bound_missing_or_quarantined_record_retains_last_good(self):
|
||||
directory = self.skill()
|
||||
self.configure(allowlist=["alpha"], always=["alpha"])
|
||||
good = self.registry.snapshot()
|
||||
(directory / "SKILL.md").write_text("broken", encoding="utf-8")
|
||||
self.assertIs(self.registry.snapshot(), good)
|
||||
self.configure(allowlist=["missing"], scenes={"help": ["missing"]})
|
||||
self.assertIs(self.registry.snapshot(), good)
|
||||
|
||||
def test_name_validation(self):
|
||||
for name in ("Upper", "two--parts", "-first", "last-", "with_space", "x" * 65, "../escape"):
|
||||
self.configure(allowlist=[name])
|
||||
self.assertFalse(self.registry.snapshot().enabled)
|
||||
self.skill("abc-123")
|
||||
self.configure(allowlist=["abc-123"])
|
||||
self.assertIn("abc-123", self.registry.snapshot().records)
|
||||
|
||||
def test_safe_reference_errors_and_no_script_exposure(self):
|
||||
self.skill(references={"note.md": "note", "script.py": "raise RuntimeError()"})
|
||||
self.configure(allowlist=["alpha"])
|
||||
snapshot = self.registry.snapshot()
|
||||
self.assertEqual(tuple(snapshot.records["alpha"].references), ("references/note.md",))
|
||||
for path in ("../SKILL.md", "/etc/passwd", "references/../SKILL.md", "references//note.md",
|
||||
"references/./note.md", "references\\note.md", "references/script.py", "note.md",
|
||||
"references/\x00.md", "references/dir/", 3):
|
||||
self.error("invalid_path", snapshot.read, "alpha", path)
|
||||
self.error("unknown_reference", snapshot.read, "alpha", "references/missing.md")
|
||||
self.error("unknown_skill", snapshot.load, "../../secret")
|
||||
|
||||
def test_binary_oversized_body_and_references(self):
|
||||
directory = self.skill()
|
||||
self.configure(allowlist=["alpha"], limits={"body_bytes": 5, "reference_bytes": 5})
|
||||
self.assertNotIn("alpha", self.registry.snapshot().records)
|
||||
self.skill(body="短文") # six UTF-8 bytes, not two
|
||||
self.assertNotIn("alpha", self.registry.snapshot().records)
|
||||
self.skill(body="ok", references={"big.txt": "123456"})
|
||||
self.assertNotIn("alpha", self.registry.snapshot().records)
|
||||
ref = directory / "references" / "big.txt"
|
||||
for data in (b"\xff", b"\x00", b"hi\x01", b"\x7f"):
|
||||
ref.write_bytes(data)
|
||||
self.assertNotIn("alpha", self.registry.snapshot().records)
|
||||
ref.write_text("ok", encoding="utf-8")
|
||||
self.assertIn("alpha", self.registry.snapshot().records)
|
||||
(directory / "SKILL.md").write_bytes(b"\x00")
|
||||
self.assertNotIn("alpha", self.registry.snapshot().records)
|
||||
|
||||
def test_symlinks_and_outside_config_are_rejected(self):
|
||||
outside = self.base / "outside"
|
||||
outside.mkdir()
|
||||
(outside / "SKILL.md").write_text("---\nname: escape\ndescription: desc\n---\nbody")
|
||||
(self.root / "escape").symlink_to(outside, target_is_directory=True)
|
||||
self.configure(allowlist=["escape"])
|
||||
self.assertFalse(self.registry.snapshot().records)
|
||||
directory = self.skill()
|
||||
refs = directory / "references"
|
||||
refs.mkdir()
|
||||
(refs / "cycle").symlink_to(refs, target_is_directory=True)
|
||||
self.configure(allowlist=["alpha"])
|
||||
self.assertFalse(self.registry.snapshot().records)
|
||||
(refs / "cycle").unlink()
|
||||
(refs / "outside.md").symlink_to(outside / "SKILL.md")
|
||||
self.assertFalse(self.registry.snapshot().records)
|
||||
(refs / "outside.md").unlink()
|
||||
good = self.registry.snapshot()
|
||||
original = self.base / "original.yml"
|
||||
original.write_text("enabled: false")
|
||||
self.config.unlink()
|
||||
self.config.symlink_to(original)
|
||||
self.assertIs(self.registry.snapshot(), good)
|
||||
# A config outside root.parent is rejected even if it is a regular file.
|
||||
nested = outside / "nested" / "skills"
|
||||
nested.mkdir(parents=True)
|
||||
other = skills.SkillRegistry(nested, original)
|
||||
self.assertFalse(other.snapshot().enabled)
|
||||
|
||||
def test_symlinked_skill_document_and_config_parent(self):
|
||||
directory = self.skill()
|
||||
doc = directory / "SKILL.md"
|
||||
target = self.base / "document.md"
|
||||
target.write_text(doc.read_text(), encoding="utf-8")
|
||||
doc.unlink()
|
||||
doc.symlink_to(target)
|
||||
self.configure(allowlist=["alpha"])
|
||||
self.assertFalse(self.registry.snapshot().records)
|
||||
link = self.base / "link"
|
||||
link.symlink_to(self.base, target_is_directory=True)
|
||||
registry = skills.SkillRegistry(link / "skills", link / "skills_config.yml")
|
||||
self.assertFalse(registry.snapshot().enabled)
|
||||
|
||||
def test_reference_only_change_changes_hash(self):
|
||||
directory = self.skill(references={"info.md": "before"})
|
||||
self.configure(allowlist=["alpha"])
|
||||
before = self.registry.snapshot()
|
||||
(directory / "references" / "info.md").write_text("after")
|
||||
after = self.registry.snapshot()
|
||||
self.assertNotEqual(before.fingerprint, after.fingerprint)
|
||||
self.assertEqual(before.read("alpha", "references/info.md"), "before")
|
||||
|
||||
def test_config_order_and_file_timestamps_do_not_affect_hash(self):
|
||||
self.skill("alpha")
|
||||
self.skill("beta")
|
||||
self.configure(allowlist=["alpha", "beta"])
|
||||
before = self.registry.snapshot()
|
||||
self.configure(allowlist=["beta", "alpha"])
|
||||
self.assertIs(before, self.registry.snapshot())
|
||||
|
||||
def test_concurrent_snapshot_reads(self):
|
||||
self.skill()
|
||||
self.configure(allowlist=["alpha"])
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
snapshots = list(pool.map(lambda _: self.registry.snapshot(), range(30)))
|
||||
self.assertTrue(all(s is snapshots[0] for s in snapshots))
|
||||
self.assertEqual(snapshots[0].load("alpha"), "Pinned instructions.\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,437 @@
|
||||
"""Transport-only skill runtime tests: no bot imports, filesystem, or network."""
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "bot-utility"))
|
||||
import skill_runtime as runtime
|
||||
|
||||
|
||||
class Snapshot:
|
||||
def __init__(self, *, enabled=True, records=None, always=(), scenes=None, limits=None):
|
||||
self.enabled = enabled
|
||||
self.records = records if records is not None else {
|
||||
"alpha": {"name": "alpha", "description": "Alpha guidance", "body": "EXACT alpha body\n第二行", "references": {"references/note.md": "EXACT reference\n不要截断"}},
|
||||
"beta": {"name": "beta", "description": "Beta guidance", "body": "EXACT beta body", "references": {}},
|
||||
"gamma": {"name": "gamma", "description": "Gamma guidance", "body": "EXACT gamma body", "references": {}},
|
||||
"delta": {"name": "delta", "description": "Delta guidance", "body": "EXACT delta body", "references": {}},
|
||||
}
|
||||
self.always = always
|
||||
self.scenes = scenes or {}
|
||||
self.limits = {"max_rounds": 2, "max_skills": 3, "max_tokens": 12000, **(limits or {})}
|
||||
self.fingerprint = "fake-pinned-fingerprint"
|
||||
self.loads = []
|
||||
self.reads = []
|
||||
|
||||
def catalog_text(self):
|
||||
return "\n".join(name + ": " + self.records[name]["description"] for name in sorted(self.records))
|
||||
|
||||
def bound_names(self, scene):
|
||||
bindings = self.scenes.get(scene, ())
|
||||
if scene == "reminder":
|
||||
bindings = self.scenes.get("notification", ()) + bindings
|
||||
return tuple(dict.fromkeys(self.always + bindings))
|
||||
|
||||
def load(self, name):
|
||||
self.loads.append(name)
|
||||
return self.records[name]["body"]
|
||||
|
||||
def read(self, name, path):
|
||||
self.reads.append((name, path))
|
||||
return self.records[name]["references"][path]
|
||||
|
||||
|
||||
def tool(action="load", **params):
|
||||
return {"name": "skill", "args": {"action": action, "params": params}}
|
||||
|
||||
|
||||
def chat(*tools, content="final"):
|
||||
return json.dumps({"message": [{"content": content}], "tools": list(tools)}, ensure_ascii=False)
|
||||
|
||||
|
||||
def text(*tools, content="final text"):
|
||||
return json.dumps({"text": content, "tools": list(tools)}, ensure_ascii=False)
|
||||
|
||||
|
||||
class SkillRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.snapshot = Snapshot()
|
||||
self.registry = Mock()
|
||||
self.registry.snapshot.return_value = self.snapshot
|
||||
self.patch = patch.object(runtime, "registry", self.registry)
|
||||
self.patch.start()
|
||||
self.addCleanup(self.patch.stop)
|
||||
print_patch = patch.object(runtime, "print", create=True)
|
||||
self.summary_print = print_patch.start()
|
||||
self.addCleanup(print_patch.stop)
|
||||
|
||||
async def test_final_prose_reuses_pinned_styles_without_json_or_tools(self):
|
||||
self.snapshot.always = ('alpha',)
|
||||
self.snapshot.scenes = {'tool_result': ('beta',)}
|
||||
with runtime.skill_session('private_chat') as turn:
|
||||
material = runtime.loaded_style_context('tool_result')
|
||||
self.assertIn('EXACT alpha body', material)
|
||||
self.assertIn('EXACT beta body', material)
|
||||
self.assertNotIn('[Skill output contract]', material)
|
||||
self.assertNotIn('[Full skill catalog]', material)
|
||||
self.assertEqual(turn.rounds, 0)
|
||||
call = AsyncMock(return_value='自然回复')
|
||||
with runtime.suspend_skills():
|
||||
result = await runtime.generate_with_skills(call, material, 'user', reasoning=False)
|
||||
self.assertEqual(result, '自然回复')
|
||||
self.assertIs(runtime.current_turn(), turn)
|
||||
self.registry.snapshot.assert_called_once()
|
||||
self.assertEqual(runtime.loaded_style_context(), '')
|
||||
|
||||
async def test_no_context_exact_identity_passthrough(self):
|
||||
system, user, images, answer = object(), object(), object(), object()
|
||||
call = AsyncMock(return_value=answer)
|
||||
result = await runtime.generate_with_skills(call, system, user, 7, presence_penalty=.7, image_infos=images, reasoning=False)
|
||||
self.assertIs(result, answer)
|
||||
call.assert_awaited_once_with(system, user, 7, presence_penalty=.7, image_infos=images, reasoning=False)
|
||||
self.registry.snapshot.assert_not_called()
|
||||
|
||||
async def test_disabled_and_empty_exact_passthrough(self):
|
||||
for snapshot in (Snapshot(enabled=False), Snapshot(records={})):
|
||||
self.registry.snapshot.return_value = snapshot
|
||||
call = AsyncMock(return_value="not JSON and unmodified")
|
||||
with runtime.skill_session("help", "text"):
|
||||
self.assertEqual(runtime.estimated_overhead(), 0)
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "system", "user"), "not JSON and unmodified")
|
||||
call.assert_awaited_once_with("system", "user", 0)
|
||||
|
||||
async def test_mixed_plan_discarded_and_transport_kwargs_preserved(self):
|
||||
business = {"name": "send_mail", "args": {"body": "PROVISIONAL SECRET"}}
|
||||
final = chat({"name": "web_search", "args": {"query": "safe"}})
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha"), business, content="PROVISIONAL"), final])
|
||||
images = [{"url": "image"}]
|
||||
with runtime.skill_session("private_chat"):
|
||||
result = await runtime.generate_with_skills(call, "ORIGINAL SYSTEM", "ORIGINAL USER", 4, image_infos=images, presence_penalty=.5, reasoning=False)
|
||||
self.assertEqual(result, final)
|
||||
self.assertEqual(call.await_count, 2)
|
||||
second = call.await_args_list[1]
|
||||
self.assertIn(self.snapshot.records["alpha"]["body"], second.args[0])
|
||||
self.assertNotIn("PROVISIONAL", second.args[0])
|
||||
self.assertIn("ORIGINAL SYSTEM", second.args[0])
|
||||
self.assertEqual(second.args[1:], ("ORIGINAL USER", 4))
|
||||
self.assertIs(second.kwargs["image_infos"], images)
|
||||
self.assertEqual(second.kwargs["presence_penalty"], .5)
|
||||
self.assertTrue(second.kwargs["reasoning"])
|
||||
|
||||
async def test_load_then_read_exact_reference_and_final_disabled(self):
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha")), chat(tool("read", name="alpha", path="references/note.md")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U"), chat())
|
||||
self.assertEqual(turn.rounds, 2)
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertIn(self.snapshot.records["alpha"]["body"], prompt)
|
||||
self.assertIn(self.snapshot.records["alpha"]["references"]["references/note.md"], prompt)
|
||||
self.assertIn("skill tools disabled", prompt)
|
||||
self.assertEqual(self.snapshot.reads, [("alpha", "references/note.md")])
|
||||
|
||||
async def test_read_requires_loaded_skill_and_invalid_path_safe(self):
|
||||
call = AsyncMock(side_effect=[chat(tool("read", name="alpha", path="references/note.md")), chat(tool(name="alpha"), tool("read", name="alpha", path="../../secrets")), chat()])
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertIn("skill_not_loaded", prompt)
|
||||
self.assertIn("invalid_reference", prompt)
|
||||
self.assertNotIn("../../secrets", prompt)
|
||||
self.assertEqual(len(self.snapshot.reads), 1)
|
||||
|
||||
async def test_duplicate_material_not_injected_twice(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha"), tool("read", name="alpha", path="references/note.md"), tool("read", name="alpha", path="references/note.md")), chat()])
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
await runtime.generate_with_skills(AsyncMock(return_value=chat()), "S", "U", reasoning=True)
|
||||
self.assertEqual(self.snapshot.loads, ["alpha"])
|
||||
self.assertEqual(len(self.snapshot.reads), 1)
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertEqual(prompt.count("EXACT alpha body"), 1)
|
||||
self.assertEqual(prompt.count("EXACT reference"), 1)
|
||||
|
||||
async def test_repeat_read_loop_fails_closed_after_one_final_call(self):
|
||||
call = AsyncMock(return_value=chat(tool(name="alpha"), {"name": "send_mail"}))
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertTrue(turn.failed)
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertEqual(call.await_count, 3)
|
||||
self.assertEqual(self.snapshot.loads, ["alpha"])
|
||||
self.assertTrue(all(item.kwargs["reasoning"] for item in call.await_args_list[1:]))
|
||||
|
||||
async def test_skill_count_includes_scene_and_always_bindings(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
self.snapshot.scenes = {"private_chat": ("beta", "gamma")}
|
||||
call = AsyncMock(side_effect=[chat(tool(name="delta")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
self.assertEqual(set(turn.loaded), {"alpha", "beta", "gamma"})
|
||||
self.assertIn("skill_budget", call.await_args.args[0])
|
||||
self.assertNotIn("EXACT delta body", call.await_args.args[0])
|
||||
|
||||
async def test_budget_rejects_whole_body_not_truncated(self):
|
||||
huge = "大" * 12500
|
||||
self.snapshot.records["alpha"]["body"] = huge
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
self.assertNotIn("alpha", turn.loaded)
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertIn('"reason": "token_budget"', prompt)
|
||||
self.assertNotIn("大", prompt)
|
||||
self.assertGreater(runtime.estimate_tokens("中" * 10), runtime.estimate_tokens("a" * 10))
|
||||
|
||||
async def test_reference_budget_rejects_whole_reference(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
self.snapshot.records["alpha"]["references"]["references/huge.md"] = "巨" * 12000
|
||||
call = AsyncMock(side_effect=[chat(tool("read", name="alpha", path="references/huge.md")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
self.assertFalse(turn.references)
|
||||
self.assertIn("token_budget", call.await_args.args[0])
|
||||
self.assertNotIn("巨", call.await_args.args[0])
|
||||
|
||||
async def test_text_envelope_after_original_and_unwrap(self):
|
||||
call = AsyncMock(side_effect=[text(tool(name="alpha"), content="NEVER SEND"), text(content="最终正文")])
|
||||
original = "只输出正文，不要JSON"
|
||||
with runtime.skill_session("help", "text"):
|
||||
self.assertEqual(await runtime.generate_with_skills(call, original, "U"), "最终正文")
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertLess(prompt.index(original), prompt.index("Internal output contract"))
|
||||
self.assertNotIn("NEVER SEND", prompt)
|
||||
|
||||
async def test_text_business_tools_rejected_including_mixed(self):
|
||||
for tools in ([{"name": "mail"}], [tool(name="alpha"), {"name": "mail"}]):
|
||||
call = AsyncMock(return_value=text(*tools))
|
||||
with runtime.skill_session("notification", "text"):
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertEqual(call.await_count, 1)
|
||||
self.assertEqual(self.snapshot.loads, [])
|
||||
|
||||
async def test_malformed_envelopes_bounded_no_raw_leak(self):
|
||||
for raw in ("PRIVATE RAW TEXT", '{"message":"not a list"}', '{"message":[],"tools":"bad"}'):
|
||||
call = AsyncMock(return_value=raw)
|
||||
with runtime.skill_session("private_chat"):
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertEqual(call.await_count, 3)
|
||||
self.assertNotIn(raw, call.await_args.args[0])
|
||||
self.assertIn("invalid_envelope", call.await_args.args[0])
|
||||
|
||||
async def test_none_immediate_no_extra_api_calls(self):
|
||||
for responses in ([None], [chat(tool(name="alpha")), None]):
|
||||
call = AsyncMock(side_effect=responses)
|
||||
with runtime.skill_session("private_chat"):
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertEqual(call.await_count, len(responses))
|
||||
|
||||
async def test_final_chat_raw_identity_and_missing_tools(self):
|
||||
final = '```json\n{"message":[{"content":"ok"}],"memory":["keep"],"need_reasoning":true}\n```'
|
||||
with runtime.skill_session("private_chat"):
|
||||
self.assertIs(await runtime.generate_with_skills(AsyncMock(return_value=final), "S", "U"), final)
|
||||
|
||||
async def test_invalid_skill_requests_safe_observations(self):
|
||||
bad = [tool("write", name="alpha", path="PRIVATE-PATH"), {"name": "skill", "args": []}, tool(name="PRIVATE-UNKNOWN"), tool("load", name=["alpha"])]
|
||||
call = AsyncMock(side_effect=[chat(*bad), chat()])
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertIn("invalid_request", prompt)
|
||||
self.assertIn("unknown_skill", prompt)
|
||||
self.assertNotIn("PRIVATE-", prompt)
|
||||
|
||||
async def test_list_uses_full_deterministic_catalog(self):
|
||||
call = AsyncMock(side_effect=[chat(tool("list")), chat()])
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
prompts = [item.args[0] for item in call.await_args_list]
|
||||
for prompt in prompts:
|
||||
self.assertEqual(prompt.count(self.snapshot.catalog_text()), 1)
|
||||
self.assertEqual(prompts[0].split("S\n\n")[0], prompts[1].split("S\n\n")[0])
|
||||
|
||||
async def test_pinned_snapshot_followups_and_nested_new_sessions(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
self.snapshot.scenes = {"tool_result": ("beta",)}
|
||||
original = self.snapshot
|
||||
replacement = Snapshot(always=("gamma",))
|
||||
with runtime.skill_session("private_chat") as outer:
|
||||
self.registry.snapshot.return_value = replacement
|
||||
with runtime.skill_scene("tool_result") as same:
|
||||
self.assertIs(same, outer)
|
||||
self.assertEqual(set(same.loaded), {"alpha", "beta"})
|
||||
self.assertIs(same.snapshot, original)
|
||||
with runtime.skill_session("help", "text") as inner:
|
||||
self.assertIsNot(inner, outer)
|
||||
self.assertIs(inner.snapshot, replacement)
|
||||
self.assertEqual(set(inner.loaded), {"gamma"})
|
||||
self.assertIs(runtime.current_turn(), outer)
|
||||
call = AsyncMock(return_value=chat())
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
self.assertIn('"scene": "private_chat"', call.await_args.args[0])
|
||||
self.assertIn("EXACT beta body", call.await_args.args[0])
|
||||
self.assertIsNone(runtime.current_turn())
|
||||
self.assertEqual(self.registry.snapshot.call_count, 2)
|
||||
|
||||
async def test_parallel_session_isolation(self):
|
||||
both_entered = asyncio.Event()
|
||||
entered = []
|
||||
async def worker(name):
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
turn.load(name)
|
||||
entered.append(name)
|
||||
if len(entered) == 2:
|
||||
both_entered.set()
|
||||
await both_entered.wait()
|
||||
self.assertIs(runtime.current_turn(), turn)
|
||||
call = AsyncMock(return_value=chat())
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
return set(turn.loaded), call.await_args.args[0]
|
||||
first, second = await asyncio.gather(worker("alpha"), worker("beta"))
|
||||
self.assertEqual(first[0], {"alpha"})
|
||||
self.assertEqual(second[0], {"beta"})
|
||||
self.assertNotIn("EXACT beta body", first[1])
|
||||
self.assertNotIn("EXACT alpha body", second[1])
|
||||
self.assertIsNone(runtime.current_turn())
|
||||
|
||||
async def test_no_implicit_session_suspend_and_exception_restoration(self):
|
||||
with runtime.skill_scene("tool_result") as absent:
|
||||
self.assertIsNone(absent)
|
||||
self.assertEqual(runtime.estimated_overhead(), 0)
|
||||
self.registry.snapshot.assert_not_called()
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
self.assertEqual(runtime.estimated_overhead(), 12000)
|
||||
with self.assertRaisesRegex(RuntimeError, "expected"):
|
||||
with runtime.suspend_skills():
|
||||
self.assertIsNone(runtime.current_turn())
|
||||
call = AsyncMock(return_value="plain")
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U"), "plain")
|
||||
call.assert_awaited_once_with("S", "U", 0)
|
||||
raise RuntimeError("expected")
|
||||
self.assertIs(runtime.current_turn(), turn)
|
||||
self.assertEqual(runtime.estimated_overhead("private_chat"), 12000)
|
||||
|
||||
async def test_followup_reuses_round_budget_and_forces_reasoning(self):
|
||||
first = AsyncMock(side_effect=[chat(tool(name="alpha")), chat()])
|
||||
followup = AsyncMock(side_effect=[chat(tool(name="beta")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
await runtime.generate_with_skills(first, "S", "U")
|
||||
with runtime.skill_scene("tool_result"):
|
||||
await runtime.generate_with_skills(followup, "S2", "U2", reasoning=False)
|
||||
self.assertEqual(turn.rounds, 2)
|
||||
self.assertTrue(followup.await_args_list[0].kwargs["reasoning"])
|
||||
self.assertEqual(followup.await_args.args[0].count("EXACT alpha body"), 1)
|
||||
|
||||
async def test_always_stable_prefix_scene_and_original_not_truncated(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
self.snapshot.scenes = {"tool_result": ("beta",)}
|
||||
original = "ORIGINAL-START" + "历史" * 20000 + "ORIGINAL-END"
|
||||
call = AsyncMock(return_value=chat())
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, original, "U")
|
||||
with runtime.skill_scene("tool_result"):
|
||||
await runtime.generate_with_skills(call, original, "U")
|
||||
first, second = [item.args[0] for item in call.await_args_list]
|
||||
self.assertIn(original, first)
|
||||
self.assertIn(original, second)
|
||||
self.assertEqual(first.split("[Skill source")[0], second.split("[Skill source")[0])
|
||||
self.assertLess(second.index("EXACT alpha body"), second.index(original))
|
||||
self.assertLess(second.index("EXACT beta body"), second.index(original))
|
||||
|
||||
async def test_reminder_inherits_notification_binding(self):
|
||||
self.snapshot.scenes = {"notification": ("alpha",)}
|
||||
with runtime.skill_session("reminder", "text") as turn:
|
||||
self.assertEqual(set(turn.loaded), {"alpha"})
|
||||
|
||||
async def test_tiny_budget_fails_before_transport(self):
|
||||
self.snapshot.limits["max_tokens"] = 10
|
||||
call = AsyncMock(return_value=chat())
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
self.assertIsNone(await runtime.generate_with_skills(call, "S", "U"))
|
||||
self.assertEqual(turn.observations[-1]["reason"], "token_budget")
|
||||
call.assert_not_awaited()
|
||||
|
||||
async def test_final_skill_disabled_allows_final_business_plan(self):
|
||||
final = chat({"name": "mail"})
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha")), chat(tool(name="beta")), final])
|
||||
with runtime.skill_session("private_chat"):
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U"), final)
|
||||
self.assertEqual(call.await_count, 3)
|
||||
self.assertIn("skill tools disabled", call.await_args.args[0])
|
||||
|
||||
async def test_text_malformed_is_not_leaked_and_can_repair(self):
|
||||
call = AsyncMock(side_effect=["PRIVATE RAW ANSWER", text(content="repaired")])
|
||||
with runtime.skill_session("help", "text"):
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U"), "repaired")
|
||||
self.assertNotIn("PRIVATE RAW ANSWER", call.await_args.args[0])
|
||||
self.assertTrue(call.await_args.kwargs["reasoning"])
|
||||
|
||||
async def test_load_exposes_reference_names_but_not_unread_source(self):
|
||||
self.snapshot.records["alpha"]["references"]["../private.md"] = "hidden"
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha")), chat()])
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "S", "U")
|
||||
prompt = call.await_args.args[0]
|
||||
self.assertIn('"available_references": ["references/note.md"]', prompt)
|
||||
self.assertNotIn("EXACT reference", prompt)
|
||||
self.assertNotIn("../private.md", prompt)
|
||||
self.assertEqual(self.snapshot.reads, [])
|
||||
|
||||
async def test_reserved_budget_never_forwarded_even_disabled(self):
|
||||
call = AsyncMock(return_value="plain")
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U", _prompt_token_budget=1), "plain")
|
||||
call.assert_awaited_once_with("S", "U", 0)
|
||||
self.registry.snapshot.assert_not_called()
|
||||
|
||||
async def test_total_text_budget_rejects_without_cutting_or_transport(self):
|
||||
call = AsyncMock(return_value=chat())
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
result = await runtime.generate_with_skills(call, "S" * 4000, "中" * 1000, _prompt_token_budget=1000)
|
||||
self.assertIsNone(result)
|
||||
self.assertEqual(turn.observations[-1]["reason"], "total_token_budget")
|
||||
call.assert_not_awaited()
|
||||
|
||||
async def test_total_text_budget_uses_engine_compatible_estimate(self):
|
||||
call = AsyncMock(return_value=chat())
|
||||
original = "a" * 20000
|
||||
with runtime.skill_session("private_chat"):
|
||||
result = await runtime.generate_with_skills(call, original, "U", _prompt_token_budget=6000)
|
||||
self.assertEqual(result, chat())
|
||||
self.assertNotIn("_prompt_token_budget", call.await_args.kwargs)
|
||||
self.assertIn(original, call.await_args.args[0])
|
||||
|
||||
async def test_small_valid_budget_can_still_load_short_source(self):
|
||||
self.snapshot.limits["max_tokens"] = 1200
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha")), chat()])
|
||||
with runtime.skill_session("private_chat") as turn:
|
||||
self.assertEqual(await runtime.generate_with_skills(call, "S", "U"), chat())
|
||||
self.assertIn("alpha", turn.loaded)
|
||||
|
||||
async def test_summary_prints_safe_metadata_and_bound_priority(self):
|
||||
self.snapshot.always = ("alpha",)
|
||||
with runtime.skill_session("private_chat"):
|
||||
pass
|
||||
summary = self.summary_print.call_args.args[0]
|
||||
self.assertIn("[skills]", summary)
|
||||
self.assertIn("fake-pinned-fingerprint", summary)
|
||||
self.assertIn('"ids": ["alpha"]', summary)
|
||||
self.assertIn('"tokens":', summary)
|
||||
self.assertNotIn("EXACT alpha body", summary)
|
||||
self.assertIn("personal expression preferences override skill", runtime._RULES)
|
||||
|
||||
async def test_logs_do_not_contain_sources_or_user_data(self):
|
||||
call = AsyncMock(side_effect=[chat(tool(name="alpha"), tool(name="PRIVATE USER SECRET")), chat()])
|
||||
with self.assertLogs("skill_runtime", "DEBUG") as logs:
|
||||
with runtime.skill_session("private_chat"):
|
||||
await runtime.generate_with_skills(call, "SECRET SYSTEM", "SECRET USER")
|
||||
output = "\n".join(logs.output)
|
||||
self.assertNotIn("SECRET", output)
|
||||
self.assertNotIn("EXACT alpha body", output)
|
||||
self.assertIn("hash=", output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Integration tests for learning and sending user-owned stickers."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import main
|
||||
from chat import engine
|
||||
|
||||
|
||||
class FakeStore:
|
||||
def __init__(self, item=None):
|
||||
self.item = item
|
||||
self.learned = []
|
||||
self.random_calls = []
|
||||
def learn(self, scope, source, classification):
|
||||
self.learned.append((scope, source, classification))
|
||||
return {"id": "learned"}
|
||||
def random_for_scene(self, scope, scene):
|
||||
self.random_calls.append((scope, scene))
|
||||
return self.item
|
||||
|
||||
|
||||
class StickerLearningTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_only_marked_stickers_are_learned(self):
|
||||
store = FakeStore()
|
||||
images = [
|
||||
{"kind": "sticker", "source_id": "s1", "original_path": "/tmp/a.webp"},
|
||||
{"kind": "photo", "source_id": "s2", "original_path": "/tmp/b.jpg"},
|
||||
]
|
||||
parsed = {"learned_stickers": [
|
||||
{"source_id": "s1", "accept": True, "scenes": ["sad"], "description": "委屈"},
|
||||
{"source_id": "s2", "accept": True, "scenes": ["happy"], "description": "普通照片"},
|
||||
]}
|
||||
with patch.object(engine, "STICKER_STORE", store), \
|
||||
patch.object(engine, "stickers_enabled", return_value=True):
|
||||
engine._learn_stickers(1001, images, parsed)
|
||||
self.assertEqual(len(store.learned), 1)
|
||||
self.assertEqual(store.learned[0][0], engine.STICKER_GLOBAL_SCOPE)
|
||||
self.assertEqual(store.learned[0][2]["scenes"], ["sad"])
|
||||
|
||||
async def test_source_marker_preserves_user_instruction(self):
|
||||
text = engine._with_sticker_sources(
|
||||
"这个适合什么时候用",
|
||||
[{"kind": "sticker", "source_id": "s1"}, {"kind": "photo"}],
|
||||
)
|
||||
self.assertIn("source_id=s1", text)
|
||||
self.assertTrue(text.endswith("这个适合什么时候用"))
|
||||
|
||||
async def test_random_candidate_requires_visual_confirmation(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
original = Path(directory, "a.webp")
|
||||
preview = Path(directory, "a.png")
|
||||
original.write_bytes(b"sticker")
|
||||
preview.write_bytes(b"preview")
|
||||
item = {"id": "a", "path": str(original), "preview_path": str(preview)}
|
||||
store = FakeStore(item)
|
||||
bot = SimpleNamespace(send_sticker=AsyncMock())
|
||||
channel = SimpleNamespace(bot=bot, _to_raw_target=lambda raw: raw)
|
||||
with patch.object(engine, "STICKER_STORE", store), \
|
||||
patch.object(engine, "stickers_enabled", return_value=True), \
|
||||
patch.object(engine, "_resolve_delivery", return_value=(channel, 10001)), \
|
||||
patch.object(engine, "_confirm_sticker", new=AsyncMock(return_value=False)):
|
||||
sent = await engine._maybe_send_sticker(channel, 1, "sad", "context")
|
||||
self.assertFalse(sent)
|
||||
self.assertEqual(store.random_calls, [(engine.STICKER_GLOBAL_SCOPE, "sad")])
|
||||
bot.send_sticker.assert_not_awaited()
|
||||
|
||||
async def test_confirmation_uses_candidate_preview(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
preview = Path(directory, "preview.png")
|
||||
preview.write_bytes(b"preview")
|
||||
model = AsyncMock(return_value='{"send":true}')
|
||||
with patch.object(engine, "call_deepseek", new=model):
|
||||
approved = await engine._confirm_sticker(
|
||||
{"preview_path": str(preview), "description": "轻微委屈"},
|
||||
"sad", "用户有点失落",
|
||||
)
|
||||
self.assertTrue(approved)
|
||||
self.assertEqual(model.await_args.kwargs["image_infos"][0]["url"], str(preview))
|
||||
|
||||
|
||||
class TelegramStickerInputTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_tg_sticker_is_downloaded_and_forwarded_as_sticker(self):
|
||||
sticker = SimpleNamespace(
|
||||
is_animated=False,
|
||||
is_video=False,
|
||||
file_unique_id="unique-1",
|
||||
emoji="🙂",
|
||||
set_name="set-1",
|
||||
)
|
||||
message = SimpleNamespace(
|
||||
from_user=SimpleNamespace(id=123, is_bot=False),
|
||||
chat=SimpleNamespace(id=123, type="private"),
|
||||
text=None,
|
||||
caption=None,
|
||||
photo=[],
|
||||
sticker=sticker,
|
||||
entities=[],
|
||||
)
|
||||
async def download(obj, destination):
|
||||
Path(destination).write_bytes(b"sticker")
|
||||
bot = SimpleNamespace(download=download, id=999)
|
||||
dispatch = AsyncMock()
|
||||
with patch.object(main, "handle_message", new=dispatch):
|
||||
await main.make_tg_handler(SimpleNamespace(tg_username="bot"))(message, bot)
|
||||
images = dispatch.await_args.args[5]
|
||||
self.assertEqual(len(images), 1)
|
||||
self.assertEqual(images[0]["kind"], "sticker")
|
||||
self.assertEqual(images[0]["source_id"], "s1")
|
||||
self.assertEqual(images[0]["file_unique_id"], "unique-1")
|
||||
self.assertEqual(images[0]["format"], "static")
|
||||
self.assertTrue(Path(images[0]["original_path"]).exists())
|
||||
Path(images[0]["original_path"]).unlink()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,260 @@
|
||||
"""Offline unit tests for the scope-isolated sticker store."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
# Locally this file and sticker_store.py share a directory. In deployment the
|
||||
# test lives in backend/tests and the module in bot-utility/chat.
|
||||
TEST_DIR = Path(__file__).resolve().parent
|
||||
BACKEND = TEST_DIR.parent
|
||||
sys.path[:0] = [
|
||||
str(TEST_DIR),
|
||||
str(BACKEND / "bot-utility" / "chat"),
|
||||
str(BACKEND / "bot-utility"),
|
||||
]
|
||||
|
||||
from sticker_store import SCENES, StickerStore, scene_prompt # noqa: E402
|
||||
|
||||
|
||||
class StickerStoreTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temporary = tempfile.TemporaryDirectory()
|
||||
self.base = Path(self.temporary.name)
|
||||
self.store = StickerStore(self.base / "store")
|
||||
self.assets = self.base / "assets"
|
||||
self.assets.mkdir()
|
||||
|
||||
def tearDown(self):
|
||||
self.temporary.cleanup()
|
||||
|
||||
def make_source(self, name="one.webp", unique_id="unique-1", body=b"sticker"):
|
||||
original = self.assets / name
|
||||
original.write_bytes(body)
|
||||
preview = self.assets / f"{original.stem}-preview.png"
|
||||
preview.write_bytes(b"preview-" + body)
|
||||
return {
|
||||
"original_path": str(original),
|
||||
"url": str(preview),
|
||||
"file_unique_id": unique_id,
|
||||
"emoji": "🙂",
|
||||
"set_name": "test-set",
|
||||
"source_uid": "sender-1",
|
||||
}
|
||||
|
||||
def test_scene_prompt_lists_all_supported_scenes(self):
|
||||
prompt = scene_prompt()
|
||||
for scene, chinese_description in SCENES.items():
|
||||
self.assertIn(scene, prompt)
|
||||
self.assertIn(chinese_description, prompt)
|
||||
|
||||
def test_learn_saves_object_preview_index_and_metadata(self):
|
||||
source = self.make_source()
|
||||
record = self.store.learn(
|
||||
"user-1",
|
||||
source,
|
||||
{"accept": True, "scenes": ["happy", "daily_cute"], "description": "开心挥手"},
|
||||
)
|
||||
|
||||
self.assertIsNotNone(record)
|
||||
self.assertEqual(record["scenes"], ["happy", "daily_cute"])
|
||||
self.assertEqual(record["description"], "开心挥手")
|
||||
self.assertEqual(record["emoji"], "🙂")
|
||||
self.assertEqual(record["set_name"], "test-set")
|
||||
self.assertEqual(record["source_uid"], "sender-1")
|
||||
self.assertEqual(record["file_unique_id"], "unique-1")
|
||||
self.assertEqual(record["format"], "webp")
|
||||
self.assertTrue(record["created_at"])
|
||||
self.assertEqual(len(record["id"]), 32)
|
||||
|
||||
object_path = Path(record["path"])
|
||||
preview_path = Path(record["preview_path"])
|
||||
self.assertTrue(object_path.is_absolute())
|
||||
self.assertTrue(preview_path.is_absolute())
|
||||
self.assertEqual(object_path.read_bytes(), b"sticker")
|
||||
self.assertEqual(preview_path.read_bytes(), b"preview-sticker")
|
||||
self.assertEqual(object_path.parent, self.base / "store" / "user-1" / "objects")
|
||||
self.assertEqual(preview_path.parent, self.base / "store" / "user-1" / "previews")
|
||||
|
||||
index_path = self.base / "store" / "user-1" / "index.json"
|
||||
saved_index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
self.assertIn(record["id"], saved_index["stickers"])
|
||||
self.assertEqual(StickerStore(self.base / "store").get("user-1", record["id"]), record)
|
||||
|
||||
def test_invalid_scenes_and_explicit_rejection_are_not_learned(self):
|
||||
invalid_classifications = [
|
||||
{"accept": False, "scenes": ["happy"]},
|
||||
{"scenes": []},
|
||||
{"scenes": ["not-a-scene"]},
|
||||
{"scenes": ["happy", "not-a-scene"]},
|
||||
{"scenes": ["happy", "sad", "angry", "thanks"]},
|
||||
{"scenes": "happy"},
|
||||
]
|
||||
for number, classification in enumerate(invalid_classifications):
|
||||
with self.subTest(classification=classification):
|
||||
source = self.make_source(
|
||||
name=f"invalid-{number}.webp", unique_id=f"invalid-{number}"
|
||||
)
|
||||
self.assertIsNone(self.store.learn("user-1", source, classification))
|
||||
self.assertEqual(self.store.prompt_summary("user-1"), "")
|
||||
|
||||
def test_random_for_scene_chooses_only_same_scene_candidates(self):
|
||||
first = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("first.webp", "first"),
|
||||
{"scenes": ["happy"], "description": "第一张"},
|
||||
)
|
||||
second = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("second.webp", "second"),
|
||||
{"scenes": ["happy", "sad"], "description": "第二张"},
|
||||
)
|
||||
self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("third.webp", "third"),
|
||||
{"scenes": ["thinking"], "description": "第三张"},
|
||||
)
|
||||
|
||||
with patch("sticker_store.random.choice", side_effect=lambda items: items[0]) as choice:
|
||||
selected = self.store.random_for_scene("user-1", "happy")
|
||||
candidates = choice.call_args.args[0]
|
||||
self.assertEqual(selected["id"], candidates[0]["id"])
|
||||
self.assertEqual({item["id"] for item in candidates}, {first["id"], second["id"]})
|
||||
|
||||
with patch("sticker_store.random.choice", side_effect=lambda items: items[-1]) as choice:
|
||||
selected = self.store.random_for_scene("user-1", "happy")
|
||||
self.assertEqual(selected["id"], choice.call_args.args[0][-1]["id"])
|
||||
self.assertIsNone(self.store.random_for_scene("user-1", "angry"))
|
||||
|
||||
def test_legacy_scopes_merge_into_shared_global_library(self):
|
||||
first = self.store.learn(
|
||||
1001, self.make_source("first.webp", "first"),
|
||||
{"scenes": ["happy"], "description": "私聊学到"},
|
||||
)
|
||||
second = self.store.learn(
|
||||
2002, self.make_source("second.webp", "second"),
|
||||
{"scenes": ["sad"], "description": "群聊学到"},
|
||||
)
|
||||
migrated = self.store.migrate_all_to_global()
|
||||
self.assertEqual(migrated, 2)
|
||||
self.assertIsNotNone(self.store.random_for_scene("global", "happy"))
|
||||
self.assertIsNotNone(self.store.random_for_scene("global", "sad"))
|
||||
self.assertIn("私聊学到", self.store.prompt_summary("global"))
|
||||
self.assertIn("群聊学到", self.store.prompt_summary("global"))
|
||||
self.assertFalse((self.base / "store" / "1001").exists())
|
||||
self.assertTrue((self.base / "store" / ".migrated" / "1001").is_dir())
|
||||
self.assertEqual(self.store.migrate_all_to_global(), 0)
|
||||
self.assertNotEqual(first["id"], second["id"])
|
||||
|
||||
def test_failed_record_prevents_scope_archive_until_retry_succeeds(self):
|
||||
valid = self.store.learn(
|
||||
1001, self.make_source("valid.webp", "valid"),
|
||||
{"scenes": ["happy"], "description": "可迁移"},
|
||||
)
|
||||
broken = self.store.learn(
|
||||
1001, self.make_source("broken.webp", "broken"),
|
||||
{"scenes": ["sad"], "description": "源文件丢失"},
|
||||
)
|
||||
broken_path = Path(broken["path"])
|
||||
broken_path.unlink()
|
||||
|
||||
self.assertEqual(self.store.migrate_all_to_global(), 1)
|
||||
self.assertTrue((self.base / "store" / "1001").is_dir())
|
||||
self.assertFalse((self.base / "store" / ".migrated" / "1001").exists())
|
||||
self.assertIsNotNone(self.store.random_for_scene("global", "happy"))
|
||||
self.assertIsNone(self.store.random_for_scene("global", "sad"))
|
||||
|
||||
broken_path.write_bytes(b"repaired")
|
||||
self.assertEqual(self.store.migrate_all_to_global(), 2)
|
||||
self.assertFalse((self.base / "store" / "1001").exists())
|
||||
self.assertTrue((self.base / "store" / ".migrated" / "1001").is_dir())
|
||||
self.assertIsNotNone(self.store.random_for_scene("global", "sad"))
|
||||
self.assertEqual(valid["file_unique_id"], "valid")
|
||||
|
||||
def test_duplicate_file_unique_id_merges_without_copying(self):
|
||||
first = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("original.webp", "same-id", b"original"),
|
||||
{"scenes": ["happy"], "description": "旧描述"},
|
||||
)
|
||||
second = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("replacement.webp", "same-id", b"replacement"),
|
||||
{"scenes": ["sad", "happy"], "description": "新描述"},
|
||||
)
|
||||
|
||||
self.assertEqual(second["id"], first["id"])
|
||||
self.assertEqual(second["scenes"], ["happy", "sad"])
|
||||
self.assertEqual(second["description"], "新描述")
|
||||
self.assertEqual(second["created_at"], first["created_at"])
|
||||
self.assertEqual(Path(second["path"]).read_bytes(), b"original")
|
||||
self.assertEqual(Path(second["preview_path"]).read_bytes(), b"preview-original")
|
||||
self.assertEqual(len(list((self.base / "store" / "user-1" / "objects").iterdir())), 1)
|
||||
self.assertEqual(len(list((self.base / "store" / "user-1" / "previews").iterdir())), 1)
|
||||
|
||||
def test_scopes_have_independent_inventory_and_deduplication(self):
|
||||
private = self.store.learn(
|
||||
"private-uid",
|
||||
self.make_source(),
|
||||
{"scenes": ["happy"], "description": "私聊贴纸"},
|
||||
)
|
||||
group = self.store.learn(
|
||||
"group-id",
|
||||
self.make_source("group.webp", body=b"group"),
|
||||
{"scenes": ["sad"], "description": "群聊贴纸"},
|
||||
)
|
||||
|
||||
self.assertNotEqual(private["id"], group["id"])
|
||||
self.assertEqual(self.store.get("private-uid", private["id"])["description"], "私聊贴纸")
|
||||
self.assertIsNone(self.store.get("private-uid", group["id"]))
|
||||
self.assertIsNone(self.store.get("group-id", private["id"]))
|
||||
self.assertIsNone(self.store.random_for_scene("private-uid", "sad"))
|
||||
self.assertEqual(self.store.random_for_scene("group-id", "sad")["id"], group["id"])
|
||||
self.assertTrue((self.base / "store" / "private-uid" / "index.json").is_file())
|
||||
self.assertTrue((self.base / "store" / "group-id" / "index.json").is_file())
|
||||
|
||||
def test_prompt_summary_contains_only_scene_counts_and_briefs(self):
|
||||
first = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("first.webp", "first"),
|
||||
{"scenes": ["happy"], "description": "开心鼓掌"},
|
||||
)
|
||||
second = self.store.learn(
|
||||
"user-1",
|
||||
self.make_source("second.webp", "second"),
|
||||
{"scenes": ["happy", "comfort"], "description": "温柔拥抱"},
|
||||
)
|
||||
|
||||
summary = self.store.prompt_summary("user-1")
|
||||
self.assertIn("happy", summary)
|
||||
self.assertIn("开心、愉快或表达喜悦", summary)
|
||||
self.assertIn("2 个", summary)
|
||||
self.assertIn("comfort", summary)
|
||||
self.assertIn("1 个", summary)
|
||||
self.assertIn("开心鼓掌", summary)
|
||||
self.assertIn("温柔拥抱", summary)
|
||||
for forbidden in (
|
||||
first["id"],
|
||||
second["id"],
|
||||
first["file_unique_id"],
|
||||
first["path"],
|
||||
first["preview_path"],
|
||||
"objects/",
|
||||
"previews/",
|
||||
):
|
||||
self.assertNotIn(forbidden, summary)
|
||||
|
||||
def test_scope_cannot_escape_root(self):
|
||||
source = self.make_source()
|
||||
for scope in ("../outside", "/tmp/outside", "a/b", "a\\b"):
|
||||
with self.subTest(scope=scope), self.assertRaises(ValueError):
|
||||
self.store.learn(scope, source, {"scenes": ["happy"]})
|
||||
self.assertFalse((self.base / "outside").exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Runtime sticker disable switch tests; no real Telegram delivery."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
from chat import engine
|
||||
|
||||
|
||||
class StickerSwitchTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_disabled_hides_tags_from_model(self):
|
||||
with patch.object(engine, "load_setting", return_value={"stickers": False}), \
|
||||
patch.object(engine.os, "listdir") as listdir:
|
||||
self.assertEqual(engine.available_sticker_tags(), [])
|
||||
listdir.assert_not_called()
|
||||
|
||||
def test_enabled_still_loads_existing_tags(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
Path(directory, "happy.webp").write_bytes(b"sticker")
|
||||
with patch.object(engine, "load_setting", return_value={"stickers": True}), \
|
||||
patch.object(engine, "STICKERS_DIR", directory):
|
||||
self.assertEqual(engine.available_sticker_tags(), ["happy"])
|
||||
|
||||
async def test_disabled_blocks_direct_send_even_with_tag(self):
|
||||
bot = type("Bot", (), {"send_sticker": AsyncMock()})()
|
||||
channel = type("Channel", (), {"bot": bot})()
|
||||
with patch.object(engine, "load_setting", return_value={"stickers": False}), \
|
||||
patch.object(engine, "_resolve_delivery") as resolve:
|
||||
await engine._maybe_send_sticker(channel, 1, "happy")
|
||||
resolve.assert_not_called()
|
||||
bot.send_sticker.assert_not_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Per-user persona overlays must be shared by chat, help, and notifications."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
|
||||
import help_ui
|
||||
import notification_render
|
||||
import notifications
|
||||
from chat import engine
|
||||
|
||||
|
||||
class UserPersonaTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_overlay_uses_long_term_preferences_with_latest_priority(self):
|
||||
rows = [
|
||||
(2, "称呼用户为主人", 5, "个人设定,表达偏好", 2, 0),
|
||||
(1, "用户不喜欢书面腔", 4, "表达偏好", 1, 0),
|
||||
]
|
||||
with patch.object(engine.memory_store, "list_facts", return_value=rows) as facts:
|
||||
prompt = engine.user_persona_prompt(42)
|
||||
facts.assert_called_once_with(42, limit=30)
|
||||
self.assertIn("当前用户的个人设定", prompt)
|
||||
self.assertIn("称呼用户为主人", prompt)
|
||||
self.assertIn("最新、最明确", prompt)
|
||||
self.assertTrue(prompt.startswith(engine.BASE_PROMPT))
|
||||
self.assertIs(engine.user_persona_prompt(None), engine.BASE_PROMPT)
|
||||
|
||||
def test_private_chat_prompt_uses_same_user_overlay(self):
|
||||
custom = engine.BASE_PROMPT + "\n【当前用户的个人设定】称呼用户为主人"
|
||||
with patch.object(engine, "user_persona_prompt", return_value=custom) as persona, \
|
||||
patch.object(engine, "load_user_status", return_value={"pause": False, "sleep": False}), \
|
||||
patch.object(engine.memory_store, "get_relation", return_value=None), \
|
||||
patch.object(engine.memory_store, "list_entities", return_value=[]), \
|
||||
patch.object(engine.memory_store, "recall_facts", return_value=[]), \
|
||||
patch.object(engine.memory_store, "migrate_legacy"), \
|
||||
patch.object(engine, "get_history_within_budget", return_value=[]), \
|
||||
patch.object(engine.notebook.notebook, "context_for_chat", return_value=""), \
|
||||
patch.object(engine, "_sticker_prompt_summary", return_value=""):
|
||||
prompt = engine.build_system_prompt(42, query="你好")
|
||||
persona.assert_called_once_with(42)
|
||||
self.assertTrue(prompt.startswith(custom))
|
||||
|
||||
async def test_notification_renderer_uses_same_user_overlay(self):
|
||||
event = notifications.Event(
|
||||
kind="komari", dedupe_key="persona-test", uid=42,
|
||||
legacy_text="raw", payload={"server": "Ama-10", "status": "offline"},
|
||||
)
|
||||
custom = engine.BASE_PROMPT + "\n【当前用户的个人设定】称呼用户为主人"
|
||||
with patch.object(engine, "user_persona_prompt", return_value=custom) as persona, \
|
||||
patch.object(engine, "call_deepseek", new=AsyncMock(return_value="主人，Ama-10 掉线了。")) as model:
|
||||
rendered = await notification_render.render_event(event, "chat", 10001)
|
||||
persona.assert_called_once_with(42)
|
||||
self.assertTrue(model.await_args.args[0].startswith(custom))
|
||||
self.assertEqual(rendered, "主人，Ama-10 掉线了。")
|
||||
|
||||
async def test_help_renderer_uses_same_user_overlay(self):
|
||||
custom = engine.BASE_PROMPT + "\n【当前用户的个人设定】称呼用户为主人"
|
||||
with patch.object(engine, "user_persona_prompt", return_value=custom) as persona, \
|
||||
patch.object(engine, "call_deepseek", new=AsyncMock(return_value="主人，可以用这些功能。")):
|
||||
await help_ui.chat_help(42)
|
||||
persona.assert_called_once_with(42)
|
||||
|
||||
def test_deterministic_notification_fallback_has_no_hardcoded_address(self):
|
||||
mail = notifications.Event(
|
||||
kind="mail", dedupe_key="mail-persona", uid=42, legacy_text="raw",
|
||||
payload={"mid": 1, "subject": "test", "sender": "a@example.test"},
|
||||
)
|
||||
ginkgo = notifications.Event(
|
||||
kind="ginkgo:course", dedupe_key="ginkgo-persona", uid=42, legacy_text="raw",
|
||||
)
|
||||
self.assertNotIn("主人", notification_render._fallback(mail))
|
||||
self.assertNotIn("主人", notification_render._fallback(ginkgo))
|
||||
|
||||
def test_memory_instruction_preserves_expression_preferences(self):
|
||||
prompt = engine.load_instruction("agent_instructions.txt")
|
||||
self.assertIn("个人设定", prompt)
|
||||
self.assertIn("表达偏好", prompt)
|
||||
self.assertIn("importance=5", prompt)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,260 @@
|
||||
"""Wardrobe tests: all persistent writes are restricted to temporary databases."""
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT / "bot-utility"))
|
||||
import wardrobe as wardrobe_module
|
||||
from wardrobe import DEFAULT_DESCRIPTION, DEFAULT_ID, HELP_TEXT, WardrobeStore, format_item
|
||||
|
||||
|
||||
def ctx(uid=1, channel="telegram", msg_type="private"):
|
||||
return SimpleNamespace(user_id=uid, channel=SimpleNamespace(name=channel),
|
||||
msg_type=msg_type, is_admin=True)
|
||||
|
||||
|
||||
def context_records(text):
|
||||
return [json.loads(line) for line in text.splitlines() if line.startswith('{')]
|
||||
|
||||
|
||||
class WardrobeTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.path = Path(self.temp.name) / "data" / "wardrobe.db"
|
||||
self.store = WardrobeStore(self.path)
|
||||
self.ctx = ctx()
|
||||
|
||||
def tearDown(self):
|
||||
self.store.close()
|
||||
self.temp.cleanup()
|
||||
|
||||
def test_import_and_constructor_are_lazy_and_environment_is_supported(self):
|
||||
self.assertFalse(self.path.exists())
|
||||
with mock.patch.dict(os.environ, {"WARDROBE_DB_PATH": str(self.path)}):
|
||||
with mock.patch.object(sqlite3, "connect", side_effect=AssertionError("import IO")):
|
||||
spec = importlib.util.spec_from_file_location("isolated_wardrobe", wardrobe_module.__file__)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
self.assertEqual(module.wardrobe.db_path, str(self.path))
|
||||
self.assertEqual(WardrobeStore().db_path, str(self.path))
|
||||
self.assertFalse(self.path.parent.exists())
|
||||
with mock.patch.dict(os.environ, {}, clear=True):
|
||||
self.assertEqual(WardrobeStore().db_path, str(ROOT / "bot-utility" / "data" / "wardrobe.db"))
|
||||
|
||||
def test_default_once_restart_and_uid_seed(self):
|
||||
records = self.store.list(self.ctx)
|
||||
self.assertEqual(len(records), 1)
|
||||
default = records[0]
|
||||
self.assertEqual(default["id"], DEFAULT_ID)
|
||||
self.assertEqual(default["name"], "白色连衣裙")
|
||||
self.assertEqual(default["category"], "连衣裙")
|
||||
self.assertEqual(default["description"], DEFAULT_DESCRIPTION)
|
||||
self.assertEqual(set(default), {"id", "uid", "name", "category", "description", "created_at", "updated_at"})
|
||||
self.store.close()
|
||||
self.assertEqual(self.store.list(self.ctx), records)
|
||||
self.assertEqual(self.store.list(ctx(2))[0]["uid"], 2)
|
||||
with sqlite3.connect(self.path) as con:
|
||||
self.assertEqual(con.execute("SELECT count(*) FROM wardrobe_users").fetchone()[0], 2)
|
||||
|
||||
def test_default_deleted_never_reseeded_including_restart(self):
|
||||
self.store.delete(self.ctx, DEFAULT_ID)
|
||||
self.assertEqual(self.store.list(self.ctx), [])
|
||||
self.store.close()
|
||||
other = WardrobeStore(self.path)
|
||||
try:
|
||||
self.assertEqual(other.list(self.ctx), [])
|
||||
text = other.context_for_chat(1)
|
||||
self.assertIn("空；总计 0 件，展示 0 件", text)
|
||||
self.assertNotIn(DEFAULT_DESCRIPTION, text)
|
||||
finally:
|
||||
other.close()
|
||||
|
||||
def test_default_update_survives_restart(self):
|
||||
updated = self.store.update(self.ctx, DEFAULT_ID, name="修改后的白裙", description="新的完整描述", category="裙子")
|
||||
self.store.close()
|
||||
self.assertEqual(self.store.list(self.ctx), [updated])
|
||||
self.assertIn("新的完整描述", self.store.context_for_chat(1))
|
||||
self.assertNotIn(DEFAULT_DESCRIPTION, self.store.context_for_chat(1))
|
||||
|
||||
def test_concurrent_first_seed_and_duplicate_across_connections(self):
|
||||
barrier = threading.Barrier(12)
|
||||
def worker(index):
|
||||
store = WardrobeStore(self.path)
|
||||
try:
|
||||
barrier.wait()
|
||||
return store.add(self.ctx, "蓝衬衫", "相同描述", "衬衫")
|
||||
finally:
|
||||
store.close()
|
||||
with ThreadPoolExecutor(max_workers=12) as pool:
|
||||
records = list(pool.map(worker, range(12)))
|
||||
self.assertTrue(all(record == records[0] for record in records))
|
||||
self.assertEqual(len(self.store.list(self.ctx)), 2)
|
||||
with sqlite3.connect(self.path) as con:
|
||||
self.assertEqual(con.execute("SELECT count(*) FROM wardrobe_users WHERE uid=1").fetchone()[0], 1)
|
||||
|
||||
def test_shared_instance_concurrent_access(self):
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
lists = list(pool.map(lambda _: self.store.list(self.ctx), range(32)))
|
||||
self.assertTrue(all(records == lists[0] for records in lists))
|
||||
self.assertEqual(len(lists[0]), 1)
|
||||
|
||||
def test_uid_isolation_and_channel_independence(self):
|
||||
record = self.store.add(ctx(1, "telegram"), "私有外套", "不可泄露")
|
||||
self.assertIn(record, self.store.list(ctx(1, "qq")))
|
||||
self.assertNotIn(record, self.store.list(ctx(2)))
|
||||
for action in (self.store.update, self.store.delete):
|
||||
with self.assertRaises(ValueError):
|
||||
action(ctx(2), record["id"])
|
||||
self.assertIn(record, self.store.list(self.ctx))
|
||||
self.assertNotIn("不可泄露", self.store.context_for_chat(2))
|
||||
self.store.update(ctx(2), DEFAULT_ID, name="用户二裙子")
|
||||
self.assertEqual(self.store.list(self.ctx, "白色连衣裙")[0]["name"], "白色连衣裙")
|
||||
|
||||
def test_nonprivate_denied_even_admin_without_database_creation(self):
|
||||
for msg_type in ("group", "guild", "", None, "PRIVATE"):
|
||||
context = ctx(msg_type=msg_type)
|
||||
for action in (lambda: self.store.list(context), lambda: self.store.add(context, "衣服"),
|
||||
lambda: self.store.update(context, DEFAULT_ID), lambda: self.store.delete(context, DEFAULT_ID)):
|
||||
with self.assertRaises(PermissionError):
|
||||
action()
|
||||
self.assertFalse(self.path.exists())
|
||||
|
||||
def test_uid_validation_all_public_methods(self):
|
||||
for uid in (True, False, 0, -1, "1", 1.0, None, 2**63):
|
||||
for action in (lambda: self.store.list(ctx(uid)), lambda: self.store.add(ctx(uid), "衣服"),
|
||||
lambda: self.store.update(ctx(uid), DEFAULT_ID), lambda: self.store.delete(ctx(uid), DEFAULT_ID),
|
||||
lambda: self.store.context_for_chat(uid)):
|
||||
with self.assertRaises(ValueError):
|
||||
action()
|
||||
self.assertFalse(self.path.exists())
|
||||
|
||||
def test_crud_full_info_idempotency_and_queries(self):
|
||||
record = self.store.add(self.ctx, " 蓝色 外套 ", " 全棉 宽松 ", " 外套 ")
|
||||
self.assertEqual(record["name"], "蓝色 外套")
|
||||
self.assertEqual(record["description"], "全棉 宽松")
|
||||
self.assertRegex(record["id"], r"^[0-9a-f]{12}$")
|
||||
self.assertEqual(record, self.store.add(self.ctx, "蓝色 外套", "全棉 宽松", "外套"))
|
||||
self.assertEqual(record, self.store.update(self.ctx, record["id"]))
|
||||
self.assertEqual(record, self.store.update(self.ctx, record["id"], name=" 蓝色 外套 ", description="全棉 宽松", category="外套"))
|
||||
self.assertEqual(self.store.list(self.ctx, "全棉"), [record])
|
||||
self.assertEqual(self.store.list(self.ctx, "外套"), [record])
|
||||
self.assertEqual(self.store.list(self.ctx, "不存在"), [])
|
||||
updated = self.store.update(self.ctx, record["id"], name="新外套", description="", category="衣服")
|
||||
self.assertEqual(updated["description"], "")
|
||||
self.assertEqual(updated["created_at"], record["created_at"])
|
||||
self.assertEqual(json.loads(format_item(updated)), updated)
|
||||
self.assertEqual(self.store.delete(self.ctx, record["id"]), updated)
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.delete(self.ctx, record["id"])
|
||||
with self.assertRaises(TypeError):
|
||||
self.store.add(self.ctx, "衣服", owner=2)
|
||||
|
||||
def test_arbitrary_categories_and_long_descriptions(self):
|
||||
for category in ("内衣", "内裤", "袜子", "首饰", "自定义分类"):
|
||||
record = self.store.add(self.ctx, category, "细" * 2000, category)
|
||||
self.assertEqual(self.store.list(self.ctx, category)[0], record)
|
||||
self.assertIn("细" * 2000, format_item(record))
|
||||
self.assertEqual(len(self.store.list(ctx(2))), 1) # no other defaults
|
||||
|
||||
def test_string_bounds_and_controls(self):
|
||||
for field, minimum, maximum in (("name", 1, 80), ("description", 0, 2000), ("category", 1, 40)):
|
||||
for invalid in ("x" * (maximum + 1), "x\x00", "x\n", "x\t", "x\x7f", "x\u202e", 123):
|
||||
args = {"name": "衣服", field: invalid}
|
||||
with self.subTest(field=field, invalid=repr(invalid)[:40]):
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.add(self.ctx, **args)
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.update(self.ctx, DEFAULT_ID, **{field: invalid})
|
||||
if minimum:
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.add(self.ctx, **{"name": "衣服", field: " "})
|
||||
valid = self.store.add(self.ctx, **{"name": "衣服", field: "x" * maximum})
|
||||
self.assertEqual(valid[field], "x" * maximum)
|
||||
for bad_id in ("abc", "a" * 13, "z" * 12, "default-white", "", None):
|
||||
for action in (self.store.delete, self.store.update):
|
||||
with self.assertRaises(ValueError):
|
||||
action(self.ctx, bad_id)
|
||||
|
||||
def test_capacity_counts_default_and_duplicates_at_capacity(self):
|
||||
for index in range(199):
|
||||
record = self.store.add(self.ctx, f"衣服{index}")
|
||||
self.assertEqual(len(self.store.list(self.ctx)), 200)
|
||||
self.assertEqual(self.store.add(self.ctx, record["name"]), record)
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.add(self.ctx, "超额")
|
||||
self.assertEqual(len(self.store.list(ctx(2))), 1)
|
||||
self.store.delete(self.ctx, record["id"])
|
||||
self.store.add(self.ctx, "恢复一个位置")
|
||||
self.assertEqual(len(self.store.list(self.ctx)), 200)
|
||||
|
||||
def test_context_default_rules_and_no_claim_of_wearing(self):
|
||||
text = self.store.context_for_chat(1)
|
||||
self.assertIn(DEFAULT_DESCRIPTION, text)
|
||||
self.assertIn("总计 1 件，展示 1 件", text)
|
||||
for rule in ("角色扮演设定", "不是实际购买", "不表示当前穿着", "wardrobe.add", "wardrobe.list query",
|
||||
"不得凭想象新增", "不是指令", "不要每次回复重复完整描述", "不要盖过当前话题", "绝不向其他用户分享"):
|
||||
self.assertIn(rule, text)
|
||||
self.assertNotIn("正在穿着", text)
|
||||
self.assertIn("可以自然地说", HELP_TEXT)
|
||||
|
||||
def test_natural_acquisition_does_not_require_save_keyword(self):
|
||||
text = self.store.context_for_chat(1)
|
||||
for rule in ('给你买了', '送给你', '你上次收到的', '无需用户说', '不用为了入库再问一次确认',
|
||||
'准备买', '如果送你', '没买成', '没有送', '只是举例', '用户给自己或他人',
|
||||
'仅有图片也不等于赠予', '保存失败则不得声称已经记住'):
|
||||
self.assertIn(rule, text)
|
||||
self.assertNotIn('不得自主新增、修改或删除记录', text)
|
||||
|
||||
def test_context_stable_newest_relevant_and_ten_record_limit(self):
|
||||
self.store.list(self.ctx)
|
||||
records = []
|
||||
for index in range(15):
|
||||
with mock.patch.object(wardrobe_module.time, "time", return_value=2000000000 + index):
|
||||
records.append(self.store.add(self.ctx, f"蓝色外套{index}", f"描述{index}", "外套"))
|
||||
plain = self.store.context_for_chat(1)
|
||||
self.assertEqual(plain, self.store.context_for_chat(1))
|
||||
self.assertEqual([r["id"] for r in context_records(plain)],
|
||||
[r["id"] for r in records[-1:-10:-1]] + [DEFAULT_ID])
|
||||
self.assertIn("总计 16 件，展示 10 件", plain)
|
||||
query = self.store.context_for_chat(1, "我想找蓝色外套0看看")
|
||||
self.assertEqual(context_records(query)[0], records[0])
|
||||
query = self.store.context_for_chat(1, "白色 连衣裙")
|
||||
self.assertEqual(context_records(query)[0]["id"], DEFAULT_ID)
|
||||
|
||||
def test_context_accepts_long_multiline_chat_query(self):
|
||||
self.store.add(self.ctx, "蓝色短裙", "底部一条白色环绕条纹", "短裙")
|
||||
text = self.store.context_for_chat(1, "聊天\n" * 1000 + "蓝色短裙")
|
||||
self.assertIn("蓝色短裙", text)
|
||||
self.assertIn(DEFAULT_DESCRIPTION, text)
|
||||
|
||||
def test_context_selects_whole_records_under_budget(self):
|
||||
for index in range(12):
|
||||
self.store.add(self.ctx, f"长描述{index}", str(index % 10) * 2000)
|
||||
text = self.store.context_for_chat(1, "长描述0")
|
||||
self.assertLessEqual(len(text), 6000)
|
||||
records = context_records(text)
|
||||
self.assertEqual(records[0]["name"], "长描述0")
|
||||
self.assertLess(len(records), 10)
|
||||
self.assertIn(f"总计 13 件，展示 {len(records)} 件", text)
|
||||
for record in records:
|
||||
if record["id"] != DEFAULT_ID:
|
||||
self.assertEqual(len(record["description"]), 2000)
|
||||
self.assertIn(record, self.store.list(self.ctx))
|
||||
# JSON escaping can make a valid record too large; omit, never truncate it.
|
||||
self.store.add(self.ctx, "引号", '"' * 2000)
|
||||
self.assertLessEqual(len(self.store.context_for_chat(1)), 6000)
|
||||
self.assertEqual(self.store.list(self.ctx, "引号")[0]["description"], '"' * 2000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Private wardrobe tool, scope and prompt integration; no network or real state."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / 'bot-utility'), str(BASE)]
|
||||
import wardrobe as module
|
||||
from agent import ToolContext, ToolExecutor
|
||||
from chat import engine
|
||||
|
||||
|
||||
class WardrobeIntegrationTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.store = module.WardrobeStore(Path(self.tmp.name) / 'wardrobe.db')
|
||||
self.channel = SimpleNamespace(name='tg', is_admin=lambda uid: False)
|
||||
self.ctx = ToolContext(self.channel, 'private', 1001, 123)
|
||||
self.other = ToolContext(self.channel, 'private', 1002, 456)
|
||||
wrapped = SimpleNamespace(wardrobe=self.store, format_item=module.format_item)
|
||||
self.executor = ToolExecutor(artifacts=MagicMock(), modules={'wardrobe': wrapped})
|
||||
|
||||
def tearDown(self):
|
||||
if hasattr(self.store, 'close'):
|
||||
self.store.close()
|
||||
self.tmp.cleanup()
|
||||
|
||||
async def test_add_and_read_full_description_and_other_owner_denied(self):
|
||||
description = '一条蓝色的、底部带一条白色环绕条纹的短裙。'
|
||||
result = await self.executor.execute('wardrobe', {'action': 'add', 'params': {
|
||||
'name': '蓝色短裙', 'category': '短裙', 'description': description}}, self.ctx)
|
||||
self.assertTrue(result.ok, result.safe_message)
|
||||
item = result.data['item']
|
||||
self.assertEqual(item['description'], description)
|
||||
read = await self.executor.execute('wardrobe', {'action': 'list', 'params': {'query': '蓝色'}}, self.ctx)
|
||||
self.assertEqual(read.data['items'][0]['description'], description)
|
||||
failed = await self.executor.execute('wardrobe', {'action': 'delete', 'params': {'id': item['id']}}, self.other)
|
||||
self.assertFalse(failed.ok)
|
||||
self.assertIn('蓝色短裙', self.store.context_for_chat(1001))
|
||||
self.assertNotIn('蓝色短裙', self.store.context_for_chat(1002))
|
||||
|
||||
async def test_tool_receipts_never_expose_record_json_or_metadata(self):
|
||||
added = await self.executor.execute('wardrobe', {'action': 'add', 'params': {
|
||||
'name': '蓝色短裙', 'description': '白色环绕条纹'}}, self.ctx)
|
||||
item_id = added.data['item']['id']
|
||||
updated = await self.executor.execute('wardrobe', {'action': 'update', 'params': {
|
||||
'id': item_id, 'description': '裙摆底部白色环绕条纹'}}, self.ctx)
|
||||
listed = await self.executor.execute('wardrobe', {'action': 'list', 'params': {}}, self.ctx)
|
||||
deleted = await self.executor.execute('wardrobe', {'action': 'delete', 'params': {'id': item_id}}, self.ctx)
|
||||
for result in (added, updated, listed, deleted):
|
||||
self.assertTrue(result.ok, result.safe_message)
|
||||
self.assertNotIn('{', result.safe_message)
|
||||
self.assertNotIn(item_id, result.safe_message)
|
||||
self.assertNotIn('uid', result.safe_message)
|
||||
for item in result.data.get('items', [result.data.get('item', {})]):
|
||||
self.assertNotIn('uid', item)
|
||||
self.assertNotIn('created_at', item)
|
||||
self.assertNotIn('updated_at', item)
|
||||
# Storage still retains full records; only outward receipts were changed.
|
||||
persisted = self.store.list(self.ctx)[0]
|
||||
self.assertIn('uid', persisted)
|
||||
|
||||
async def test_group_cannot_list_or_mutate_even_as_admin(self):
|
||||
group = ToolContext(self.channel, 'group', 1001, 999, is_admin=True)
|
||||
self.assertNotIn('wardrobe', [x['name'] for x in self.executor.prompt_schema(group)])
|
||||
self.assertIn('wardrobe', [x['name'] for x in self.executor.prompt_schema(self.ctx)])
|
||||
for action, params in [('list', {}), ('add', {'name': '袜子'}),
|
||||
('update', {'id': 'default-white-dress', 'name': 'x'}),
|
||||
('delete', {'id': 'default-white-dress'})]:
|
||||
result = await self.executor.execute('wardrobe', {'action': action, 'params': params}, group)
|
||||
self.assertEqual(result.code, 'permission_denied')
|
||||
|
||||
async def test_model_cannot_override_owner(self):
|
||||
result = await self.executor.execute('wardrobe', {'action': 'add', 'params': {
|
||||
'name': '袜子', 'uid': 1002}}, self.ctx)
|
||||
self.assertEqual(result.code, 'invalid_args')
|
||||
|
||||
async def test_list_pages_preserve_whole_records(self):
|
||||
for n in range(12):
|
||||
self.store.add(self.ctx, f'衣物{n}', '细节' * 900, '其他')
|
||||
first = await self.executor.execute('wardrobe', {'action': 'list', 'params': {'limit': 20}}, self.ctx)
|
||||
self.assertTrue(first.ok)
|
||||
self.assertGreater(first.data['next_offset'], 0)
|
||||
self.assertLess(first.data['next_offset'], first.data['total'])
|
||||
for item in first.data['items']:
|
||||
if item['id'] != 'default-white-dress': self.assertEqual(item['description'], '细节' * 900)
|
||||
second = await self.executor.execute('wardrobe', {'action': 'list', 'params': {
|
||||
'offset': first.data['next_offset']}}, self.ctx)
|
||||
ids = {x['id'] for x in first.data['items']}
|
||||
self.assertFalse(ids.intersection(x['id'] for x in second.data['items']))
|
||||
|
||||
def test_private_and_proactive_prompt_include_wardrobe_after_stable_instructions(self):
|
||||
with patch.object(engine, 'user_persona_prompt', return_value='PERSONA'), \
|
||||
patch.object(engine, 'load_user_status', return_value={'pause': False, 'sleep': False}), \
|
||||
patch.object(engine.memory_store, 'get_relation', return_value=None), \
|
||||
patch.object(engine.memory_store, 'list_entities', return_value=[]), \
|
||||
patch.object(engine.memory_store, 'recall_facts', return_value=[]), \
|
||||
patch.object(engine.memory_store, 'migrate_legacy'), \
|
||||
patch.object(engine, 'get_history_within_budget', return_value=[]), \
|
||||
patch.object(engine, 'get_last_active', return_value=0), \
|
||||
patch.object(engine, '_sticker_prompt_summary', return_value=''), \
|
||||
patch.object(engine.notebook.notebook, 'context_for_chat', return_value=''), \
|
||||
patch.object(engine.wardrobe, 'wardrobe', self.store):
|
||||
for is_auto in (False, True):
|
||||
prompt = engine.build_system_prompt(1001, is_auto=is_auto)
|
||||
self.assertIn('白色连衣裙', prompt)
|
||||
self.assertLess(prompt.index('当你回复时'), prompt.index('白色连衣裙'))
|
||||
self.assertNotIn('现在是 ', prompt)
|
||||
|
||||
async def test_group_prompt_never_reads_private_wardrobe(self):
|
||||
with patch.object(engine, 'get_recent_group_history', return_value=[]), \
|
||||
patch.object(engine, 'call_deepseek', new=AsyncMock(return_value=json.dumps({'message': [], 'tools': []}))), \
|
||||
patch.object(engine.wardrobe.wardrobe, 'context_for_chat') as context:
|
||||
await engine._call_group_model_and_reply_impl(None, 999, '你好', 0)
|
||||
context.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Wardrobe orchestration regressions: fake model, executor and task store only."""
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
BASE = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
|
||||
import agent_loop
|
||||
from agent import ToolResult
|
||||
|
||||
|
||||
def call(name="白衬衫", description="柔软", action="add", tool="wardrobe"):
|
||||
return {"name": tool, "args": {"action": action, "params": {"name": name, "description": description}}}
|
||||
|
||||
|
||||
def request(*calls):
|
||||
return {"message": [], "tools": list(calls)}
|
||||
|
||||
|
||||
def reply(text):
|
||||
return {"message": [{"content": text}], "tools": []}
|
||||
|
||||
|
||||
class Channel:
|
||||
name = "test"
|
||||
|
||||
def is_admin(self, uid):
|
||||
return False
|
||||
|
||||
|
||||
class Store:
|
||||
last = None
|
||||
|
||||
def __init__(self):
|
||||
self.steps = []
|
||||
self.finished = None
|
||||
Store.last = self
|
||||
|
||||
def start(self, *args):
|
||||
return "fake-task"
|
||||
|
||||
def record(self, *args):
|
||||
self.steps.append(args)
|
||||
|
||||
def finish(self, task, status, error=""):
|
||||
self.finished = status
|
||||
|
||||
|
||||
class Engine:
|
||||
BASE_PROMPT = "GROUP_PERSONA"
|
||||
user_persona_prompt = staticmethod(lambda uid: "ORIGINAL_PERSONA")
|
||||
|
||||
def __init__(self, responses):
|
||||
self.responses = list(responses)
|
||||
self.prompts = []
|
||||
|
||||
async def call_deepseek(self, prompt, *args, **kwargs):
|
||||
polishing = "【仅润色回复正文】" in prompt
|
||||
assert kwargs["reasoning"] is (not polishing)
|
||||
self.prompts.append(prompt)
|
||||
response = self.responses.pop(0) if self.responses else RuntimeError("outage")
|
||||
if isinstance(response, Exception):
|
||||
raise response
|
||||
return json.dumps(response, ensure_ascii=False) if polishing and isinstance(response, dict) else response
|
||||
|
||||
def parse_json_response(self, raw):
|
||||
return raw
|
||||
|
||||
|
||||
class WardrobeReplyTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def run_loop(self, initial, responses, outcomes=None, **limits):
|
||||
calls = []
|
||||
outcomes = list(outcomes or [])
|
||||
|
||||
class Executor:
|
||||
async def execute(self, name, args, ctx):
|
||||
calls.append((name, args))
|
||||
if outcomes:
|
||||
return outcomes.pop(0)
|
||||
item = {"id": len(calls), "uid": 999, "name": args["params"]["name"],
|
||||
"description": "RAW_PRIVATE_DESCRIPTION", "created_at": "SECRET_TIME"}
|
||||
return ToolResult(True, "ok", json.dumps(item), data={"item": item})
|
||||
|
||||
engine = Engine(responses)
|
||||
with patch.object(agent_loop, "AgentTaskStore", Store), patch.object(agent_loop, "ToolExecutor", Executor):
|
||||
with patch.multiple(agent_loop, **({"MAX_CALLS": 12, "MAX_ROUNDS": 8} | limits)):
|
||||
result = await agent_loop.run_agent_loop(engine, initial, "ORIGINAL_PERSONA", "user", Channel(), "private", 1, 2)
|
||||
return result, calls, engine
|
||||
|
||||
async def test_plain_prose_finalization_does_not_need_json_fields(self):
|
||||
result, calls, model = await self.run_loop(request(call()), [None, "谢谢主人，这件我记住啦。"])
|
||||
self.assertEqual(agent_loop._message_text(result), "谢谢主人，这件我记住啦。")
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(result['tools'], [])
|
||||
self.assertNotIn('请仅生成有效JSON', model.prompts[-1])
|
||||
self.assertNotIn('RAW_PRIVATE_DESCRIPTION', model.prompts[-1])
|
||||
self.assertNotIn('SECRET_TIME', model.prompts[-1])
|
||||
self.assertIn('ORIGINAL_PERSONA', model.prompts[-1])
|
||||
self.assertEqual(Store.last.finished, 'done')
|
||||
|
||||
async def test_legacy_text_envelope_can_omit_tools(self):
|
||||
result, calls, _ = await self.run_loop(request(call()), [None, '{"text":"谢谢主人。"}'])
|
||||
self.assertEqual(agent_loop._message_text(result), '谢谢主人。')
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
def test_plain_reply_rejects_tools_and_db_dump_without_execution(self):
|
||||
for raw in ('{"tools":[{"name":"wardrobe"}],"text":"x"}',
|
||||
'{"id":"abc","uid":42,"name":"dress"}', '', None):
|
||||
self.assertIsNone(agent_loop._plain_reply(raw))
|
||||
self.assertEqual(agent_loop._plain_reply('好呀，谢谢你。'), '好呀，谢谢你。')
|
||||
|
||||
async def test_outage_never_exposes_raw_records(self):
|
||||
result, calls, engine = await self.run_loop(request(call()), [RuntimeError("offline")])
|
||||
self.assertEqual(agent_loop._message_text(result), "好呀，白衬衫已经在衣柜里啦。")
|
||||
self.assertEqual(result["tools"], [])
|
||||
self.assertEqual(len(engine.prompts), 2)
|
||||
self.assertIn("ORIGINAL_PERSONA", engine.prompts[-1])
|
||||
self.assertIn("实际工具结果", engine.prompts[-1])
|
||||
|
||||
async def test_same_normalized_name_paraphrase_replays(self):
|
||||
result, calls, engine = await self.run_loop(
|
||||
request(call("Ａ Shirt", "丝绸")),
|
||||
[request(call("a\tshirt", "轻柔丝质")), reply("记住啦")])
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertIn('"code": "already_saved"', engine.prompts[-1])
|
||||
self.assertIn("update", engine.prompts[-1])
|
||||
self.assertIn('"id": 1', engine.prompts[-1])
|
||||
|
||||
async def test_outage_receipt_names_only_completed_adds(self):
|
||||
result, calls, _ = await self.run_loop(request(call("红裙"), call("蓝裙"), call("未执行绿裙")), [None, None])
|
||||
text = agent_loop._message_text(result)
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertEqual("好呀，红裙、蓝裙已经在衣柜里啦。", text)
|
||||
for forbidden in ("未执行绿裙", "RAW_PRIVATE_DESCRIPTION", "uid", "SECRET_TIME", "{"):
|
||||
self.assertNotIn(forbidden, text)
|
||||
|
||||
async def test_different_names_both_saved(self):
|
||||
_, calls, _ = await self.run_loop(request(call("红裙"), call("蓝裙")), [reply("好呀")])
|
||||
self.assertEqual(len(calls), 2)
|
||||
|
||||
async def test_update_reversion_is_not_replayed_from_stale_cache(self):
|
||||
first = call('外套', '蓝色', action='update')
|
||||
second = call('外套', '红色', action='update')
|
||||
_, calls, _ = await self.run_loop(request(first), [request(second), request(first), reply('改好了')])
|
||||
self.assertEqual(len(calls), 3)
|
||||
|
||||
async def test_delete_invalidates_same_turn_add_guard(self):
|
||||
added = ToolResult(True, 'ok', '已保存', data={'item': {'id': 'first', 'name': '外套'}})
|
||||
removed = ToolResult(True, 'ok', '已删除', data={'item': {'id': 'first', 'name': '外套'}})
|
||||
readded = ToolResult(True, 'ok', '已保存', data={'item': {'id': 'second', 'name': '外套'}})
|
||||
_, calls, _ = await self.run_loop(request(call('外套')),
|
||||
[request(call('外套', action='delete')), request(call('外套')), reply('收好了')],
|
||||
outcomes=[added, removed, readded])
|
||||
self.assertEqual(len(calls), 3)
|
||||
|
||||
async def test_exact_duplicates_do_not_execute(self):
|
||||
_, calls, _ = await self.run_loop(request(call(), call()), [reply("好呀")])
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
async def test_failed_add_can_retry(self):
|
||||
failed = ToolResult(False, "busy", "暂时忙，请重试")
|
||||
_, calls, _ = await self.run_loop(request(call()), [request(call()), reply("好了")], [failed])
|
||||
self.assertEqual(len(calls), 2)
|
||||
|
||||
async def test_all_prior_successes_survive_in_bounded_ledger(self):
|
||||
rounds = [request(call(f"衣服{i}")) for i in range(1, 5)]
|
||||
_, _, engine = await self.run_loop(request(call("衣服0")), rounds + [reply("好呀")])
|
||||
last = engine.prompts[-1]
|
||||
ledger = json.loads(last.split("【累计成功操作摘要（内部）】\n")[1].split("\n【你刚才")[0])
|
||||
self.assertEqual(len(ledger), 5)
|
||||
for i in range(5):
|
||||
self.assertIn(f"衣服{i}", " ".join(ledger))
|
||||
self.assertTrue(all(len(entry) <= 400 for entry in ledger))
|
||||
self.assertNotIn("RAW_PRIVATE_DESCRIPTION", " ".join(ledger))
|
||||
self.assertNotIn("SECRET_TIME", " ".join(ledger))
|
||||
self.assertIn("tool=wardrobe; action=add; code=ok", ledger[0])
|
||||
self.assertIn('"action": "add"', last)
|
||||
|
||||
async def test_ledger_bound_and_no_raw_arguments(self):
|
||||
rounds = [request(call(f"衣服{i}")) for i in range(1, 14)]
|
||||
_, _, engine = await self.run_loop(request(call("衣服0", "password=secret")), rounds + [reply("好了")], MAX_ROUNDS=20, MAX_CALLS=20)
|
||||
ledger = json.loads(engine.prompts[-1].split("【累计成功操作摘要（内部）】\n")[1].split("\n【你刚才")[0])
|
||||
self.assertEqual(len(ledger), 12)
|
||||
self.assertNotIn("secret", " ".join(ledger))
|
||||
self.assertTrue(all(len(entry) <= 400 for entry in ledger))
|
||||
|
||||
async def test_read_only_queries_are_not_deduplicated(self):
|
||||
for tool in ("wardrobe", "mail"):
|
||||
_, calls, _ = await self.run_loop(request(call(action="list", tool=tool)),
|
||||
[request(call(action="list", tool=tool)), reply("好了")])
|
||||
self.assertEqual(len(calls), 2)
|
||||
|
||||
async def test_invalid_followup_and_invalid_final_content(self):
|
||||
result, calls, engine = await self.run_loop(request(call()),
|
||||
[{"tools": "invalid", "message": []}, {"tools": [], "message": [{"content": {"text": "wrong type"}}]}])
|
||||
self.assertEqual(agent_loop._message_text(result), "好呀，白衬衫已经在衣柜里啦。")
|
||||
self.assertEqual(len(engine.prompts), 2)
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
async def test_latest_observation_keeps_all_results(self):
|
||||
_, _, engine = await self.run_loop(request(call("一"), call("二")), [reply("好了")])
|
||||
observation = engine.prompts[0].split("【你刚才执行了工具，以下是结构化结果】\n")[1].split("\n请根据")[0]
|
||||
self.assertEqual(len(json.loads(observation)), 2)
|
||||
|
||||
async def test_invalid_model_fallback_status_is_not_done(self):
|
||||
await self.run_loop(request(call()), [None, None])
|
||||
self.assertEqual(Store.last.finished, "response_fallback")
|
||||
|
||||
async def test_finalization_cannot_execute_proposed_tool(self):
|
||||
result, calls, engine = await self.run_loop(request(call()), [None, request(call("不该保存"))])
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(result["tools"], [])
|
||||
self.assertEqual(len(engine.prompts), 2)
|
||||
|
||||
async def test_leaked_final_response_gets_one_repair(self):
|
||||
result, calls, engine = await self.run_loop(request(call()), [reply('{"id":1,"description":"private"}'), reply("谢谢你，我记住啦。")])
|
||||
self.assertEqual(agent_loop._message_text(result), "谢谢你，我记住啦。")
|
||||
self.assertEqual(len(engine.prompts), 2)
|
||||
|
||||
async def test_leaked_repair_rejected(self):
|
||||
result, _, _ = await self.run_loop(request(call()), [None, reply('{"uid":1,"created_at":"today"}')])
|
||||
self.assertNotIn("uid", str(result))
|
||||
self.assertIn("已经在衣柜里", str(result))
|
||||
|
||||
async def test_empty_message_and_limits_finalize(self):
|
||||
for limits, responses in [({}, [reply(""), reply("记住了")]),
|
||||
({"MAX_ROUNDS": 1}, [request(call("后续")), reply("记住了")]),
|
||||
({"MAX_CALLS": 1}, [request(call("后续")), reply("记住了")])]:
|
||||
result, calls, engine = await self.run_loop(request(call()), responses, **limits)
|
||||
self.assertEqual(agent_loop._message_text(result), "记住了")
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(len(engine.prompts), 2)
|
||||
|
||||
async def test_mixed_failure_not_all_success(self):
|
||||
good = ToolResult(True, "ok", '{"id":1,"uid":999}', data={"item": {"name": "裙子"}})
|
||||
bad = ToolResult(False, "invalid_args", "衣物名称不能为空")
|
||||
result, _, _ = await self.run_loop(request(call("一"), call("二")), [None], [good, bad])
|
||||
text = agent_loop._message_text(result)
|
||||
self.assertIn("还有操作没完成", text)
|
||||
self.assertIn("衣物名称不能为空", text)
|
||||
self.assertNotIn("这些衣服我记住啦", text)
|
||||
self.assertNotIn("uid", text)
|
||||
|
||||
def test_list_names_only_and_write_fallbacks(self):
|
||||
record = {"id": 4, "uid": 42, "name": "绿裙", "description": "SECRET", "created_at": "DATE"}
|
||||
result = agent_loop._fallback_parsed([{"tool": "wardrobe", "action": "list", "ok": True, "safe_message": json.dumps(record), "data": {"items": [record]}}])
|
||||
self.assertEqual(agent_loop._message_text(result), "衣柜里有：绿裙。")
|
||||
for action, word in [("update", "改好"), ("delete", "移除")]:
|
||||
text = agent_loop._message_text(agent_loop._fallback_parsed([{"tool": "wardrobe", "action": action, "ok": True, "safe_message": json.dumps(record)}]))
|
||||
self.assertIn(word, text)
|
||||
self.assertNotIn("SECRET", text)
|
||||
|
||||
def test_nonwardrobe_fallback_unchanged(self):
|
||||
result = agent_loop._fallback_parsed([{"tool": "mail", "ok": True, "safe_message": "邮件已发送"}])
|
||||
self.assertEqual(agent_loop._message_text(result), "已经处理好了呢。\n邮件已发送")
|
||||
result = agent_loop._fallback_parsed([{"tool": "pixiv", "ok": False, "code": "send_failed", "safe_message": "发送失败 https://example.test"}])
|
||||
self.assertEqual(agent_loop._message_text(result), "处理时遇到了具体问题：\npixiv（send_failed）：发送失败 https://example.test")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user