from __future__ import annotations import json from datetime import UTC, datetime, timedelta from types import SimpleNamespace import pytest from aiohttp import CookieJar from aiohttp.test_utils import TestClient, TestServer def connection_payload(connection_id: str = "conn-a", *, can_reply: bool = True) -> dict: return { "id": connection_id, "user": {"id": 900, "username": "owner", "first_name": "店主"}, "user_chat_id": 900, "date": int(datetime.now(UTC).timestamp()), "rights": {"can_reply": can_reply, "can_read_messages": True}, "is_enabled": True, } def customer_message( *, message_id: int, text: str | None = "营业时间是什么？", chat_id: int = 501, sender_id: int = 501, sent_at: datetime | None = None, **values, ) -> dict: message = { "business_connection_id": "conn-a", "message_id": message_id, "date": int((sent_at or datetime.now(UTC)).timestamp()), "chat": {"id": chat_id, "type": "private"}, "from": {"id": sender_id, "username": f"user{sender_id}", "first_name": "客户"}, **values, } if text is not None: message["text"] = text return message class FakeProvider: configured = True def __init__(self, *, error: Exception | None = None) -> None: self.error = error self.calls: list[dict] = [] async def decide(self, **values): self.calls.append(values) if self.error: raise self.error entry_id = str(values["knowledge"][0]["entry_id"]) return { "action": "answer", "reply": "每天 09:00 至 18:00 营业。", "handoff_reason": "", "matched_entry_ids": [entry_id], "summary": "客户询问营业时间。", } async def test_connection(self): return {"ok": True, "model": "test-model", "response_id": "response-1"} class FakeBusinessApi: def __init__(self) -> None: self.sent: list[dict] = [] self.read: list[tuple[str, int, int]] = [] async def send_message( self, chat_id: int, text: str, *, business_connection_id: str = "", reply_markup: dict | None = None, ) -> dict: self.sent.append( { "chat_id": chat_id, "text": text, "business_connection_id": business_connection_id, "reply_markup": reply_markup, } ) return { "message_id": 1000 + len(self.sent), "sender_business_bot": {"id": 999}, } async def read_business_message( self, connection_id: str, chat_id: int, message_id: int ) -> None: self.read.append((connection_id, chat_id, message_id)) async def prepare_runtime(app_modules, *, provider: FakeProvider | None = None): dbassistant = app_modules.load("wbb.utils.dbassistant") service = app_modules.load("wbb.services.business_assistant") await dbassistant.upsert_business_connection(connection_payload()) await dbassistant.update_account_settings("conn-a", {"assistant_enabled": True}) knowledge = await dbassistant.create_knowledge_entry( "conn-a", { "question": "营业时间是什么？", "aliases": ["几点营业"], "keywords": ["营业时间", "营业"], "answer": "每天 09:00 至 18:00 营业。", "priority": 10, }, ) selected_provider = provider or FakeProvider() runtime = service.BusinessAssistantRuntime( token="123456:" + "A" * 30, session=SimpleNamespace(), provider=selected_provider, ) runtime.api = FakeBusinessApi() return dbassistant, service, runtime, knowledge async def test_connection_settings_knowledge_quota_and_bot_scope(app_modules): dbassistant = app_modules.load("wbb.utils.dbassistant") await dbassistant.upsert_business_connection(connection_payload("conn-a")) await dbassistant.upsert_business_connection(connection_payload("conn-b")) settings = await dbassistant.update_account_settings( "conn-a", { "assistant_enabled": True, "account_daily_limit": 2, "customer_daily_limit": 1, "timezone": "Asia/Shanghai", }, ) first = await dbassistant.create_knowledge_entry( "conn-a", { "question": "营业时间", "aliases": ["几点开门"], "keywords": ["开门"], "answer": "九点开门。", "priority": 20, }, ) await dbassistant.create_knowledge_entry( "conn-b", {"question": "营业时间", "keywords": ["开门"], "answer": "十点开门。"}, ) await app_modules.wbb.db.business_assistant_knowledge.insert_one( { "bot_id": "foreign-bot", "entry_id": "foreign-entry", "connection_id": "conn-a", "question": "营业时间", "aliases": [], "keywords": ["开门"], "answer": "不应跨 Bot 返回。", "priority": 1000, "enabled": True, "updated_at": datetime.now(UTC), } ) matched = await dbassistant.match_knowledge("conn-a", "你们几点开门？") assert [item["entry_id"] for item in matched] == [first["entry_id"]] items, total = await dbassistant.list_knowledge_entries("conn-a") assert total == 1 assert items[0]["answer"] == "九点开门。" await app_modules.wbb.db.business_assistant_connections.insert_one( { "bot_id": "foreign-bot", "connection_id": "foreign-connection", "updated_at": datetime.now(UTC), } ) connections, total = await dbassistant.list_business_connections(page_size=100) assert total == 2 assert {item["connection_id"] for item in connections} == {"conn-a", "conn-b"} assert await dbassistant.reserve_ai_usage("conn-a", 101, settings) == (True, "") assert await dbassistant.reserve_ai_usage("conn-a", 101, settings) == ( False, "customer_daily_limit", ) assert await dbassistant.reserve_ai_usage("conn-a", 102, settings) == (True, "") assert await dbassistant.reserve_ai_usage("conn-a", 103, settings) == ( False, "account_daily_limit", ) usage = await dbassistant.usage_metrics(connection_id="conn-a") assert usage["ai_calls"] == 2 assert usage["customers"] == 2 with pytest.raises(dbassistant.AssistantDataError) as error: dbassistant.normalize_account_settings({"timezone": "invalid/timezone"}) assert error.value.code == "invalid_setting" def test_ai_contract_and_reply_window_validation(app_modules): service = app_modules.load("wbb.services.business_assistant") parsed = service.parse_ai_decision( json.dumps( { "action": "answer", "reply": "标准答案", "matched_entry_ids": ["entry-1", "foreign"], "summary": "摘要", } ), allowed_entry_ids={"entry-1"}, ) assert parsed["matched_entry_ids"] == ["entry-1"] with pytest.raises(service.AssistantProviderError): service.parse_ai_decision( '{"action":"answer","reply":"没有引用"}', allowed_entry_ids={"entry-1"} ) now = datetime.now(UTC) assert service.business_reply_window_open( {"date": int((now - timedelta(hours=23)).timestamp())}, now=now ) assert not service.business_reply_window_open( {"date": int((now - timedelta(hours=25)).timestamp())}, now=now ) assert not service.business_reply_window_open({}, now=now) async def test_business_updates_are_idempotent_and_manual_reply_pauses(app_modules): dbassistant, _service, runtime, knowledge = await prepare_runtime(app_modules) update = { "update_id": 11, "business_message": customer_message(message_id=1), } await runtime.process_update(update) await runtime.process_update(update) assert len(runtime.api.sent) == 1 assert runtime.api.sent[0]["business_connection_id"] == "conn-a" assert runtime.api.read == [("conn-a", 501, 1)] assert len(runtime.provider.calls) == 1 assert runtime.provider.calls[0]["knowledge"][0]["entry_id"] == knowledge["entry_id"] conversation = await dbassistant.get_conversation_by_chat("conn-a", 501) assert conversation["summary"] == "客户询问营业时间。" messages = await dbassistant.recent_conversation_messages( conversation["conversation_id"], limit=10 ) assert [item["direction"] for item in messages] == ["incoming", "assistant"] assert messages[0]["expires_at"] - messages[0]["created_at"] == timedelta(days=30) message_indexes = await app_modules.wbb.db.business_assistant_messages.index_information() assert message_indexes["expires_at_1"]["expireAfterSeconds"] == 0 await runtime.process_update( { "update_id": 12, "business_message": customer_message( message_id=2, text="账号本人回复", sender_id=900 ), } ) paused = await dbassistant.get_conversation(conversation["conversation_id"]) assert paused["status"] == "human_paused" assert paused["handoff_reason"] == "human_reply_detected" assert paused["customer"]["id"] == 501 await runtime.process_update( { "update_id": 13, "business_message": customer_message( message_id=3, sender_business_bot={"id": 999}, ), } ) await runtime.process_update( { "update_id": 14, "business_message": customer_message( message_id=4, is_from_offline=True, ), } ) assert len(runtime.api.sent) == 1 async def test_handoff_sensitive_non_text_provider_error_and_expired_window(app_modules): dbassistant, service, runtime, _knowledge = await prepare_runtime(app_modules) await dbassistant.update_account_settings( "conn-a", { "notification_destination": "both", "ops_group_id": -100123, "account_daily_limit": 1, "customer_daily_limit": 1, }, ) await runtime.process_update( { "update_id": 21, "business_message": customer_message( message_id=1, text="我要退款", chat_id=601, sender_id=601 ), } ) sensitive = await dbassistant.get_conversation_by_chat("conn-a", 601) assert sensitive["status"] == "handoff" assert sensitive["handoff_reason"] == "sensitive_request" assert runtime.api.sent[-1]["business_connection_id"] == "" assert "message_id=1" in runtime.api.sent[-1]["text"] assert {item["chat_id"] for item in runtime.api.sent[-2:]} == {900, -100123} await runtime.process_update( { "update_id": 22, "business_message": customer_message( message_id=2, text=None, chat_id=602, sender_id=602, photo=[{"file_id": "x"}] ), } ) unsupported = await dbassistant.get_conversation_by_chat("conn-a", 602) assert unsupported["handoff_reason"] == "unsupported_message" runtime.provider = FakeProvider(error=service.AssistantProviderError("模型超时")) await runtime.process_update( { "update_id": 23, "business_message": customer_message( message_id=3, chat_id=603, sender_id=603 ), } ) provider_failed = await dbassistant.get_conversation_by_chat("conn-a", 603) assert provider_failed["handoff_reason"] == "provider_error" await runtime.process_update( { "update_id": 230, "business_message": customer_message( message_id=30, chat_id=605, sender_id=605 ), } ) quota_handoff = await dbassistant.get_conversation_by_chat("conn-a", 605) assert quota_handoff["handoff_reason"] == "account_daily_limit" sent_before = len(runtime.api.sent) await runtime.process_update( { "update_id": 24, "business_message": customer_message( message_id=4, chat_id=604, sender_id=604, sent_at=datetime.now(UTC) - timedelta(hours=25), ), } ) expired = await dbassistant.get_conversation_by_chat("conn-a", 604) assert expired["handoff_reason"] == "reply_window_expired" new_sends = runtime.api.sent[sent_before:] assert len(new_sends) == 2 assert all(item["business_connection_id"] == "" for item in new_sends) async def test_failed_update_goes_to_dead_letter_without_blocking(app_modules, monkeypatch): dbassistant, _service, runtime, _knowledge = await prepare_runtime(app_modules) async def fail(_message): raise RuntimeError("persistent failure") async def no_sleep(_seconds): return None runtime._handle_business_message = fail monkeypatch.setattr("wbb.services.business_assistant.asyncio.sleep", no_sleep) await runtime.process_update( {"update_id": 31, "business_message": customer_message(message_id=1)} ) dead_letter = await app_modules.wbb.db.business_assistant_dead_letters.find_one( {"bot_id": "primary", "update_id": 31} ) update = await app_modules.wbb.db.business_assistant_updates.find_one( {"bot_id": "primary", "update_id": 31} ) assert dead_letter["error"] == "persistent failure" assert update["status"] == "done" assert await dbassistant.claim_update(31) is False async def test_connection_updates_revoke_permissions_and_persist_offset(app_modules): service = app_modules.load("wbb.services.business_assistant") dbassistant = app_modules.load("wbb.utils.dbassistant") runtime = service.BusinessAssistantRuntime( token="123456:" + "A" * 30, session=SimpleNamespace(), provider=FakeProvider(), ) runtime.api = FakeBusinessApi() await runtime.process_update( {"update_id": 41, "business_connection": connection_payload(can_reply=True)} ) connected = await dbassistant.get_business_connection("conn-a") assert connected["is_enabled"] is True assert connected["rights"]["can_reply"] is True await dbassistant.update_account_settings("conn-a", {"assistant_enabled": True}) revoked = connection_payload(can_reply=False) revoked["is_enabled"] = False await runtime.process_update({"update_id": 42, "business_connection": revoked}) disconnected = await dbassistant.get_business_connection("conn-a") assert disconnected["is_enabled"] is False assert "can_reply" not in disconnected["rights"] await runtime.process_update( {"update_id": 420, "business_message": customer_message(message_id=20)} ) assert runtime.api.sent == [] await dbassistant.save_update_offset(43) assert await dbassistant.load_update_offset() == 43 await dbassistant.save_update_offset(44) assert await dbassistant.load_update_offset() == 44 class StartupApi: def __init__(self, *, supported: bool, webhook_url: str = "") -> None: self.supported = supported self.webhook_url = webhook_url async def get_me(self): return {"username": "business_bot", "can_connect_to_business": self.supported} async def get_webhook_info(self): return {"url": self.webhook_url} async def test_runtime_blocks_webhook_conflict_and_disabled_business_mode(app_modules): service = app_modules.load("wbb.services.business_assistant") runtime = service.BusinessAssistantRuntime( token="123456:" + "A" * 30, session=SimpleNamespace(), provider=FakeProvider(), ) runtime.api = StartupApi(supported=True, webhook_url="https://hooks.example.com/tg") webhook_status = await runtime.start() assert webhook_status["polling_state"] == "blocked" assert webhook_status["webhook_conflict"] is True assert runtime._poll_task is None runtime.api = StartupApi(supported=False) business_status = await runtime.start() assert business_status["polling_state"] == "blocked" assert business_status["business_mode_supported"] is False assert "BotFather" in business_status["last_error"] class FakeResponse: def __init__(self, status: int, payload: dict) -> None: self.status = status self.payload = payload async def __aenter__(self): return self async def __aexit__(self, *_args): return None async def json(self, **_kwargs): return self.payload class RetrySession: def __init__(self, responses: list[FakeResponse]) -> None: self.responses = responses self.calls = 0 def post(self, *_args, **_kwargs): response = self.responses[self.calls] self.calls += 1 return response async def test_telegram_api_retries_429_and_5xx(app_modules, monkeypatch): service = app_modules.load("wbb.services.business_assistant") session = RetrySession( [ FakeResponse( 429, { "ok": False, "error_code": 429, "description": "Too Many Requests", "parameters": {"retry_after": 3}, }, ), FakeResponse( 502, {"ok": False, "error_code": 502, "description": "Bad Gateway"}, ), FakeResponse(200, {"ok": True, "result": {"id": 999}}), ] ) sleeps: list[int] = [] async def record_sleep(seconds): sleeps.append(seconds) monkeypatch.setattr("wbb.services.business_assistant.asyncio.sleep", record_sleep) api = service.TelegramBusinessApi("123456:" + "A" * 30, session) assert await api.get_me() == {"id": 999} assert session.calls == 3 assert sleeps == [3, 2] async def _login_and_change_password(client: TestClient) -> str: login = await client.post( "/api/admin/v1/auth/login", json={"username": "admin", "password": "qwe0.123456"}, ) login_data = (await login.json())["data"] changed = await client.put( "/api/admin/v1/auth/password", headers={"X-CSRF-Token": login_data["csrf_token"]}, json={ "current_password": "qwe0.123456", "new_password": "changed-pass-123", }, ) return (await changed.json())["data"]["csrf_token"] async def test_business_assistant_api_permission_csrf_audit_and_clear(app_modules): dbassistant = app_modules.load("wbb.utils.dbassistant") await dbassistant.upsert_business_connection(connection_payload()) conversation = await dbassistant.get_or_create_conversation( "conn-a", 701, customer={"id": 701, "first_name": "待清理客户"} ) await dbassistant.append_conversation_message( conversation["conversation_id"], direction="incoming", telegram_message_id=1, text="请清理", ) await dbassistant.update_conversation_summary(conversation["conversation_id"], "待清理摘要") admin_api = app_modules.load("wbb.admin.api") application = admin_api.build_admin_application() await application["admin_api"].initialize() client = TestClient(TestServer(application), cookie_jar=CookieJar(unsafe=True)) await client.start_server() try: csrf = await _login_and_change_password(client) app_modules.wbb.BOT_PERMISSIONS = set() denied = await client.get("/api/admin/v1/business-assistant/status") assert denied.status == 403 app_modules.wbb.BOT_PERMISSIONS = {"business_assistant.manage"} no_csrf = await client.put( "/api/admin/v1/business-assistant/settings", json={"connection_id": "conn-a", "assistant_enabled": True, "confirm": True}, ) assert no_csrf.status == 403 unconfirmed = await client.put( "/api/admin/v1/business-assistant/settings", headers={"X-CSRF-Token": csrf}, json={"connection_id": "conn-a", "assistant_enabled": True}, ) assert unconfirmed.status == 409 saved = await client.put( "/api/admin/v1/business-assistant/settings", headers={"X-CSRF-Token": csrf}, json={"connection_id": "conn-a", "assistant_enabled": True, "confirm": True}, ) assert saved.status == 200 assert (await saved.json())["data"]["assistant_enabled"] is True created = await client.post( "/api/admin/v1/business-assistant/knowledge", headers={"X-CSRF-Token": csrf}, json={ "connection_id": "conn-a", "question": "配送范围", "keywords": ["配送"], "answer": "仅限市区。", "confirm": True, }, ) assert created.status == 201 listed = await client.get( "/api/admin/v1/business-assistant/knowledge?connection_id=conn-a" ) assert (await listed.json())["data"]["total"] == 1 cleared = await client.post( f"/api/admin/v1/business-assistant/conversations/{conversation['conversation_id']}/actions", headers={"X-CSRF-Token": csrf}, json={"action": "clear", "confirm": True}, ) assert cleared.status == 200 detail = await dbassistant.conversation_detail(conversation["conversation_id"]) assert detail["summary"] == "" assert detail["messages"] == [] audit = await app_modules.wbb.db.admin_audit_logs.find_one( {"action": "business_assistant.conversation.clear"} ) assert audit["success"] is True finally: await client.close()