Steven
f6e7616756
CI / backend (push) Canceled after 0s
上一版用 `manager.polling()` 判"有没有人在收消息"，那是进程内状态：请求落到持有者 回答 true、落到备用进程回答 false，页面上就会每两秒闪一次"已停止"。 改成查数据库里的 advisory lock（`pg_locks`，objsubid=2 的两段式锁），谁抢到租约所有 进程都看得见。`lease_held()` 自己开一条短会话查，不复用请求会话——接口里常在 `await session.commit()` 之后才拼返回值，那时原事务已经关了。
394 lines
14 KiB
Python
394 lines
14 KiB
Python
"""微信 clawbot：扫码登录 + 长轮询 + 消息落成任务/知识。
|
|
|
|
不联网：iLink 的四个接口都用 monkeypatch 顶掉，验的是我们这侧的流程与落库：
|
|
取二维码 → 等扫码 → 存加密登录态 → 轮询拿到消息 → 建事项 → 留痕 → 回复 → 去重。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from httpx import AsyncClient
|
|
from sqlalchemy import select, text
|
|
|
|
from app.db import SessionFactory, bind_tenant
|
|
from app.models_integrations import BotChannel, BotMessage
|
|
from app.models_task import Item
|
|
from app.services import wechat_bot, wechat_ilink
|
|
from app.services.wechat_intake import handle_text
|
|
|
|
from .conftest import create_workspace, login, unwrap
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def no_background_tasks(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""别在测试里真起轮询任务：这里只验"该不该起"，任务体另测。"""
|
|
|
|
async def fake_start(self, channel_id: uuid.UUID) -> None: # noqa: ARG001
|
|
return None
|
|
|
|
async def fake_stop(self, channel_id: uuid.UUID) -> None: # noqa: ARG001
|
|
return None
|
|
|
|
monkeypatch.setattr(wechat_bot.BotManager, "start", fake_start)
|
|
monkeypatch.setattr(wechat_bot.BotManager, "stop", fake_stop)
|
|
|
|
|
|
async def _begin(client: AsyncClient, user: dict[str, Any], ws: dict[str, Any]) -> dict[str, Any]:
|
|
return unwrap(
|
|
await client.post(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
|
|
|
|
def _patch_qr(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
async def fake_qr() -> wechat_ilink.QrCode:
|
|
return wechat_ilink.QrCode(qrcode_id="qr-1", payload="https://wx.example/qr/1")
|
|
|
|
monkeypatch.setattr(wechat_ilink, "get_qrcode", fake_qr)
|
|
|
|
|
|
def _patch_status(monkeypatch: pytest.MonkeyPatch, status: wechat_ilink.QrStatus) -> None:
|
|
async def fake_status(qrcode_id: str) -> wechat_ilink.QrStatus: # noqa: ARG001
|
|
return status
|
|
|
|
monkeypatch.setattr(wechat_ilink, "qrcode_status", fake_status)
|
|
|
|
|
|
async def test_begin_login_returns_scannable_qr(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_patch_qr(monkeypatch)
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "扫码接入")
|
|
|
|
data = await _begin(client, user, ws)
|
|
|
|
assert data["qrcode_id"] == "qr-1"
|
|
assert data["payload"] == "https://wx.example/qr/1"
|
|
assert data["image"].startswith("data:image/svg+xml;base64,")
|
|
|
|
channels = unwrap(
|
|
await client.get(f"/api/v1/workspaces/{ws['id']}/integrations", headers=user["headers"])
|
|
)
|
|
assert channels["items"][0]["status"] == "pending"
|
|
assert channels["items"][0]["connected"] is False
|
|
|
|
|
|
async def test_scan_confirms_and_stores_encrypted_token(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_patch_qr(monkeypatch)
|
|
_patch_status(
|
|
monkeypatch,
|
|
wechat_ilink.QrStatus(
|
|
status="confirmed", bot_token="bot-token-abc", ilink_user_id="wx-user-9"
|
|
),
|
|
)
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "扫码接入")
|
|
started = await _begin(client, user, ws)
|
|
|
|
result = unwrap(
|
|
await client.get(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr/{started['qrcode_id']}",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
assert result["status"] == "confirmed"
|
|
assert result["channel"]["connected"] is True
|
|
assert result["channel"]["ilink_user_id"] == "wx-user-9"
|
|
|
|
channel_id = uuid.UUID(started["channel_id"])
|
|
async with SessionFactory() as session:
|
|
await session.execute(
|
|
text("SELECT set_config('app.bot_channel', :c, true)"), {"c": str(channel_id)}
|
|
)
|
|
channel = await session.get(BotChannel, channel_id)
|
|
assert channel is not None
|
|
assert channel.bot_token_enc is not None
|
|
# 落库的是密文，且能解回原值
|
|
assert "bot-token-abc" not in channel.bot_token_enc
|
|
assert wechat_ilink.decrypt_token(channel.bot_token_enc) == "bot-token-abc"
|
|
|
|
|
|
async def test_stale_qr_poll_is_reported_as_expired(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_patch_qr(monkeypatch)
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "扫码接入")
|
|
await _begin(client, user, ws)
|
|
|
|
data = unwrap(
|
|
await client.get(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr/another-qr",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
assert data["status"] == "expired"
|
|
|
|
|
|
async def test_disconnect_clears_login_state(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_patch_qr(monkeypatch)
|
|
_patch_status(monkeypatch, wechat_ilink.QrStatus(status="confirmed", bot_token="bot-token-abc"))
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "扫码接入")
|
|
started = await _begin(client, user, ws)
|
|
unwrap(
|
|
await client.get(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr/{started['qrcode_id']}",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
|
|
data = unwrap(
|
|
await client.delete(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/{started['channel_id']}",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
assert data["channel"]["connected"] is False
|
|
assert data["channel"]["status"] == "pending"
|
|
|
|
|
|
async def test_plain_text_becomes_task(client: AsyncClient) -> None:
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "收消息")
|
|
ws_id = uuid.UUID(ws["id"])
|
|
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
channel = BotChannel(
|
|
workspace_id=ws_id, kind="wechat_clawbot", name="微信", status="active"
|
|
)
|
|
session.add(channel)
|
|
await session.flush()
|
|
result = await handle_text(
|
|
session, channel, "给客户回电话\n顺便确认报价", base_url="http://site"
|
|
)
|
|
await session.commit()
|
|
|
|
assert result["action"] == "task"
|
|
# 回话里只确认"收下了"，不带站点链接（站点地址在手机上不一定打得开）
|
|
assert result["reply"] == "已保存任务：给客户回电话"
|
|
assert "http" not in result["reply"]
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
item = await session.get(Item, result["item_id"])
|
|
assert item is not None
|
|
assert item.title == "给客户回电话"
|
|
assert item.kind == "task"
|
|
assert item.status == "inbox"
|
|
|
|
|
|
async def test_knowledge_prefix_creates_draft_with_tags(client: AsyncClient) -> None:
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "收消息")
|
|
ws_id = uuid.UUID(ws["id"])
|
|
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
channel = BotChannel(
|
|
workspace_id=ws_id, kind="wechat_clawbot", name="微信", status="active"
|
|
)
|
|
session.add(channel)
|
|
await session.flush()
|
|
result = await handle_text(
|
|
session,
|
|
channel,
|
|
"知识：回滚前先停灰度 #发布 #回滚\n先停流量，再退镜像版本",
|
|
base_url="http://site",
|
|
)
|
|
await session.commit()
|
|
|
|
assert result["action"] == "knowledge"
|
|
assert result["reply"] == "已保存知识草稿：回滚前先停灰度"
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
item = await session.get(Item, result["item_id"])
|
|
assert item is not None
|
|
assert item.kind == "knowledge"
|
|
assert item.status == "draft"
|
|
assert item.title == "回滚前先停灰度"
|
|
assert "#" not in item.title
|
|
|
|
|
|
async def test_help_does_not_create_anything(client: AsyncClient) -> None:
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "收消息")
|
|
ws_id = uuid.UUID(ws["id"])
|
|
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
channel = BotChannel(
|
|
workspace_id=ws_id, kind="wechat_clawbot", name="微信", status="active"
|
|
)
|
|
session.add(channel)
|
|
await session.flush()
|
|
result = await handle_text(session, channel, "帮助", base_url="http://site")
|
|
await session.commit()
|
|
|
|
assert result["action"] == "help"
|
|
assert result["item_id"] is None
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
items = (
|
|
(await session.execute(select(Item).where(Item.workspace_id == ws_id))).scalars().all()
|
|
)
|
|
assert items == []
|
|
|
|
|
|
async def test_poller_handles_message_once_and_replies(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""同一个外部消息 id 到两次，只能有一条任务、一条留痕、一次回复。"""
|
|
_patch_qr(monkeypatch)
|
|
_patch_status(monkeypatch, wechat_ilink.QrStatus(status="confirmed", bot_token="bot-token-abc"))
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "轮询")
|
|
started = await _begin(client, user, ws)
|
|
unwrap(
|
|
await client.get(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr/{started['qrcode_id']}",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
channel_id = uuid.UUID(started["channel_id"])
|
|
|
|
replies: list[tuple[str, str]] = []
|
|
|
|
async def fake_send(
|
|
token: str, to: str, context: str | None, text: str, *, base_url: str | None = None
|
|
) -> None: # noqa: ARG001
|
|
replies.append((to, text))
|
|
|
|
monkeypatch.setattr(wechat_ilink, "send_text", fake_send)
|
|
|
|
message = wechat_ilink.Inbound(
|
|
message_id="msg-1",
|
|
sender="wx-user-9",
|
|
text="整理一下本周排期",
|
|
context_token="ctx-1",
|
|
)
|
|
await wechat_bot.manager._handle(channel_id, message, "bot-token-abc", None)
|
|
await wechat_bot.manager._handle(channel_id, message, "bot-token-abc", None)
|
|
|
|
assert len(replies) == 1
|
|
assert replies[0][0] == "wx-user-9"
|
|
assert replies[0][1] == "已保存任务：整理一下本周排期"
|
|
|
|
async with SessionFactory() as session:
|
|
await session.execute(
|
|
text("SELECT set_config('app.bot_channel', :c, true)"), {"c": str(channel_id)}
|
|
)
|
|
channel = await session.get(BotChannel, channel_id)
|
|
assert channel is not None
|
|
assert channel.context_token == "ctx-1"
|
|
assert channel.last_seen_at is not None
|
|
# bot_message 走标准租户隔离，验它得先绑定工作区
|
|
await bind_tenant(session, uuid.UUID(ws["id"]), uuid.UUID(user["user_id"]))
|
|
messages = (
|
|
(await session.execute(select(BotMessage).where(BotMessage.channel_id == channel_id)))
|
|
.scalars()
|
|
.all()
|
|
)
|
|
assert len(messages) == 1
|
|
assert messages[0].action == "task"
|
|
|
|
|
|
async def test_stale_token_marks_channel_expired(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_patch_qr(monkeypatch)
|
|
_patch_status(monkeypatch, wechat_ilink.QrStatus(status="confirmed", bot_token="bot-token-abc"))
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "轮询")
|
|
started = await _begin(client, user, ws)
|
|
unwrap(
|
|
await client.get(
|
|
f"/api/v1/workspaces/{ws['id']}/integrations/wechat/qr/{started['qrcode_id']}",
|
|
headers=user["headers"],
|
|
)
|
|
)
|
|
channel_id = uuid.UUID(started["channel_id"])
|
|
|
|
await wechat_bot.manager._mark_relogin(channel_id, "登录态已失效，需要重新扫码")
|
|
|
|
channels = unwrap(
|
|
await client.get(f"/api/v1/workspaces/{ws['id']}/integrations", headers=user["headers"])
|
|
)
|
|
assert channels["items"][0]["status"] == "expired"
|
|
assert channels["items"][0]["connected"] is False
|
|
assert "重新扫码" in channels["items"][0]["last_error"]
|
|
|
|
|
|
async def test_upstream_wait_status_is_normalized(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""上游说 "wait"，页面只认 waiting——这一层要对齐，否则界面停在"等扫码结果…"。"""
|
|
|
|
async def fake_post(path: str, body: dict, **kwargs: Any) -> dict: # noqa: ARG001
|
|
return {"status": "wait"}
|
|
|
|
monkeypatch.setattr(wechat_ilink, "_post", fake_post)
|
|
result = await wechat_ilink.qrcode_status("qr-x")
|
|
assert result.status == "waiting"
|
|
|
|
|
|
async def test_only_one_process_holds_the_polling_lease() -> None:
|
|
"""生产是 gunicorn `-w 2`：同一个通道只能有一个进程在拉消息，另一个等着接班。"""
|
|
channel_id = uuid.uuid4()
|
|
|
|
first = await wechat_bot.acquire_lease(channel_id)
|
|
assert first is not None, "第一个进程该抢到"
|
|
|
|
second = await wechat_bot.acquire_lease(channel_id)
|
|
assert second is None, "第二个进程不该同时抢到，否则两条轮询会互相覆盖游标"
|
|
|
|
await first.close() # 进程退出，连接断，锁释放
|
|
third = await wechat_bot.acquire_lease(channel_id)
|
|
assert third is not None, "前一个进程放手后，备用进程要能接上"
|
|
await third.close()
|
|
|
|
|
|
async def test_startup_sees_logged_in_channels(client: AsyncClient) -> None:
|
|
"""重启后要能把已登录的通道接回来：没有工作区上下文时也得看得见（迁移 0016）。"""
|
|
user = await login(client)
|
|
ws = await create_workspace(client, user, "重启接管")
|
|
ws_id = uuid.UUID(ws["id"])
|
|
|
|
async with SessionFactory() as session:
|
|
await bind_tenant(session, ws_id, uuid.UUID(user["user_id"]))
|
|
channel = BotChannel(
|
|
workspace_id=ws_id,
|
|
kind="wechat_clawbot",
|
|
name="微信",
|
|
status="active",
|
|
bot_token_enc=wechat_ilink.encrypt_token("bot-token-x"),
|
|
)
|
|
session.add(channel)
|
|
await session.commit()
|
|
channel_id = channel.id
|
|
|
|
ids = await wechat_bot.active_channel_ids()
|
|
assert channel_id in ids, "启动时看不见已登录通道，重启后机器人就静默失联了"
|
|
|
|
|
|
async def test_polling_flag_is_visible_to_every_process() -> None:
|
|
"""页面的"接收中"要跨进程一致：看数据库里的锁，而不是本进程的任务表。"""
|
|
channel_id = uuid.uuid4()
|
|
assert await wechat_bot.lease_held(channel_id) is False
|
|
|
|
lease = await wechat_bot.acquire_lease(channel_id)
|
|
assert lease is not None
|
|
try:
|
|
assert await wechat_bot.lease_held(channel_id) is True
|
|
finally:
|
|
await lease.close()
|
|
assert await wechat_bot.lease_held(channel_id) is False
|