Ctrl K
ComunidadGuías…

tests/test_dingtalk_gate.py

Rama main · réplica de lectura del repositorio original.

Archivo original ↗Preguntar a KiBot
"""钉钉组织级接入的确定性单测:触发过滤 / 角色映射热加载 / 并发闸门 / LLM 流式解析。

不连钉钉、不连网关:消息用鸭子类型桩(should_trigger/resolve_identity 全走 getattr),
LLM 流式用打桩 requests.post 喂伪 SSE。真实字段名与端到端行为由 GB10 真实入口验收覆盖。
"""
import asyncio
import os
import time
from types import SimpleNamespace

import pytest

from dm.channels.dingtalk import DEFAULT_ROLE, _Gate, _RoleMap, resolve_identity, should_trigger


def _msg(text="库存多少", ctype="2", at=False, bot_id="bot1", at_ids=(), nick="张三", staff="s001"):
    return SimpleNamespace(
        text=SimpleNamespace(content=text),
        conversation_type=ctype,
        is_in_at_list=at,
        chatbot_user_id=bot_id,
        at_users=[SimpleNamespace(dingtalk_id=i) for i in at_ids],
        sender_nick=nick,
        sender_staff_id=staff,
    )


# ---------- 触发过滤 ----------

def test_group_without_at_ignored():
    ok, text = should_trigger(_msg(ctype="2", at=False))
    assert not ok and text == "库存多少"


def test_group_with_at_triggers():
    ok, text = should_trigger(_msg(ctype="2", at=True))
    assert ok and text == "库存多少"


def test_group_at_fallback_via_at_users():
    ok, _ = should_trigger(_msg(ctype="2", at=False, bot_id="bot1", at_ids=("x", "bot1")))
    assert ok


def test_private_chat_passes_through():
    ok, text = should_trigger(_msg(ctype="1", at=False))
    assert ok and text == "库存多少"


def test_empty_and_at_prefix_stripped():
    ok, _ = should_trigger(_msg(text="   ", ctype="1"))
    assert not ok
    ok, text = should_trigger(_msg(text="@数据机器人 M0001 库存", ctype="1"))
    assert ok and text == "M0001 库存"


# ---------- 角色映射(热加载) ----------

def test_role_map_missing_file_defaults(tmp_path):
    rm = _RoleMap(path=str(tmp_path / "nope.yaml"))
    assert rm.role_of("anyone") == DEFAULT_ROLE


def test_role_map_hit_flat_and_extended(tmp_path):
    p = tmp_path / "role_map.yaml"
    p.write_text('roles:\n  "s001": 采购\n  "s002": { role: 管理层, warehouse_id: W01 }\n', encoding="utf-8")
    rm = _RoleMap(path=str(p))
    assert rm.role_of("s001") == "采购"
    assert rm.role_of("s002") == "管理层"
    assert rm.role_of("s999") == DEFAULT_ROLE


def test_role_map_hot_reload_on_mtime(tmp_path):
    p = tmp_path / "role_map.yaml"
    p.write_text('roles:\n  "s001": 采购\n', encoding="utf-8")
    rm = _RoleMap(path=str(p))
    assert rm.role_of("s001") == "采购"
    p.write_text('roles:\n  "s001": 管理层\n', encoding="utf-8")
    os.utime(p, (time.time() + 5, time.time() + 5))  # 强制 mtime 前进,规避文件系统粒度
    assert rm.role_of("s001") == "管理层"


def test_role_map_hot_reload_detects_atomic_replacement_with_preserved_mtime(tmp_path):
    p = tmp_path / "role_map.yaml"
    p.write_text('roles:\n  "s001": 采购\n', encoding="utf-8")
    original = p.stat()
    rm = _RoleMap(path=str(p))
    assert rm.role_of("s001") == "采购"

    replacement = tmp_path / "replacement.yaml"
    replacement.write_text('roles:\n  "s001": 仓管\n', encoding="utf-8")
    os.utime(replacement, ns=(original.st_atime_ns, original.st_mtime_ns))
    os.replace(replacement, p)

    assert rm.role_of("s001") == "仓管"


def test_role_map_refuses_symbolic_link_targets(tmp_path):
    target = tmp_path / "outside.yaml"
    target.write_text('roles:\n  "s001": 管理员\n', encoding="utf-8")
    link = tmp_path / "role_map.yaml"
    link.symlink_to(target)

    rm = _RoleMap(path=str(link))

    assert rm.role_of("s001") == DEFAULT_ROLE


def test_role_map_broken_yaml_degrades(tmp_path):
    p = tmp_path / "role_map.yaml"
    p.write_text("roles: [::broken", encoding="utf-8")
    rm = _RoleMap(path=str(p))
    assert rm.role_of("s001") == DEFAULT_ROLE


def test_role_map_duplicate_staff_keys_degrade_instead_of_overriding(tmp_path):
    p = tmp_path / "role_map.yaml"
    p.write_text('roles:\n  "s001": 采购\n  "s001": 管理层\n', encoding="utf-8")
    rm = _RoleMap(path=str(p))
    assert rm.role_of("s001") == DEFAULT_ROLE


def test_resolve_identity_nick_then_staff_then_anon(tmp_path):
    rm = _RoleMap(path=str(tmp_path / "nope.yaml"))
    assert resolve_identity(_msg(nick="张三", staff="s1"), rm) == ("张三", DEFAULT_ROLE)
    assert resolve_identity(_msg(nick="", staff="s1"), rm)[0] == "s1"
    assert resolve_identity(_msg(nick="", staff=""), rm)[0] == "钉钉用户"


# ---------- 并发闸门 ----------

def test_gate_queue_feedback_and_timeout():
    async def main():
        gate = _Gate(limit=1, queue_max=2, queue_timeout=0.15)
        events = []

        async def job(tag, hold):
            if gate.waiting >= gate.queue_max:
                events.append((tag, "rejected_full"))
                return
            if gate.full:
                events.append((tag, "queued", gate.eta_text()))
            gate.waiting += 1
            try:
                await asyncio.wait_for(gate.sem.acquire(), timeout=gate.queue_timeout)
            except asyncio.TimeoutError:
                events.append((tag, "rejected_timeout"))
                return
            finally:
                gate.waiting -= 1
            gate.active += 1
            t0 = time.monotonic()
            try:
                events.append((tag, "run"))
                await asyncio.sleep(hold)
            finally:
                gate.active -= 1
                gate.durations.append(time.monotonic() - t0)
                gate.sem.release()

        t1 = asyncio.create_task(job("A", 0.3))   # 占住唯一槽 0.3s
        while gate.active == 0:                   # Windows 循环 sleep(0.01)≈0,须显式等 A 真拿到槽
            await asyncio.sleep(0.01)
        t2 = asyncio.create_task(job("B", 0.05))  # 满载 → 收排队反馈,0.15s 内等不到 → 超时拒绝
        await asyncio.gather(t1, t2)
        return events, gate

    events, gate = asyncio.run(main())
    tags = [(e[0], e[1]) for e in events]
    assert ("A", "run") in tags
    assert ("B", "queued") in tags
    assert ("B", "rejected_timeout") in tags
    eta_msg = next(e[2] for e in events if e[1] == "queued")
    assert "前面还有 1 个" in eta_msg and "秒" in eta_msg
    assert gate.active == 0 and gate.waiting == 0        # 终态归零
    assert len(gate.durations) == 1                       # 只有 A 真正跑完喂了耗时


def test_gate_eta_defaults_and_rounding():
    gate = _Gate(limit=4, queue_max=12, queue_timeout=300)
    assert "约 60 秒" in gate.eta_text()                  # 无样本按 60s
    gate.durations.extend([100.0])
    gate.active = 4
    assert "约 120 秒" in gate.eta_text()                 # ceil(1/4)*100 → 30s 粒度上取整


# ---------- dm.llm.chat 流式解析 ----------

class _FakeResp:
    def __init__(self, lines, status=200):
        self.status_code = status
        self._lines = lines
        self.closed = False
        self.text = ""

    def iter_lines(self, decode_unicode=True):
        yield from self._lines

    def close(self):
        self.closed = True


def _sse(*chunks, done=True):
    lines = ['data: {"choices":[{"delta":{"content":"%s"}}]}' % c for c in chunks]
    if done:
        lines.append("data: [DONE]")
    return lines


def test_llm_chat_stream_collects_content(monkeypatch):
    import dm.llm as llm
    monkeypatch.setattr(llm, "LLM_STREAMING", True)
    resp = _FakeResp(_sse("你好", ",", "世界"))
    monkeypatch.setattr(llm.requests, "post", lambda *a, **k: resp)
    assert llm.chat([{"role": "user", "content": "hi"}]) == "你好,世界"
    assert resp.closed


def test_llm_chat_stream_error_block(monkeypatch):
    import dm.llm as llm
    monkeypatch.setattr(llm, "LLM_STREAMING", True)
    resp = _FakeResp(['data: {"error":{"message":"boom"}}'])
    monkeypatch.setattr(llm.requests, "post", lambda *a, **k: resp)
    with pytest.raises(RuntimeError, match="流式响应报错"):
        llm.chat([{"role": "user", "content": "hi"}])
    assert resp.closed


def test_llm_chat_stream_wall_timeout(monkeypatch):
    import dm.llm as llm
    monkeypatch.setattr(llm, "LLM_STREAMING", True)

    class _SlowResp(_FakeResp):
        def iter_lines(self, decode_unicode=True):
            while True:  # 无限流,靠墙钟斩断
                yield 'data: {"choices":[{"delta":{"content":"x"}}]}'

    resp = _SlowResp([])
    monkeypatch.setattr(llm.requests, "post", lambda *a, **k: resp)
    t = {"n": 0}
    real = time.monotonic
    monkeypatch.setattr(llm.time, "monotonic", lambda: real() + (t.__setitem__("n", t["n"] + 1) or t["n"] * 3.0))
    with pytest.raises(RuntimeError, match="超墙钟"):
        llm.chat([{"role": "user", "content": "hi"}], timeout=5)
    assert resp.closed  # 超时后连接已关闭(取消传导)


def test_llm_chat_nonstream_fallback(monkeypatch):
    import dm.llm as llm
    monkeypatch.setattr(llm, "LLM_STREAMING", False)

    class _JsonResp:
        status_code = 200
        text = ""

        def json(self):
            return {"choices": [{"message": {"content": "老路径"}}]}

    monkeypatch.setattr(llm.requests, "post", lambda *a, **k: _JsonResp())
    assert llm.chat([{"role": "user", "content": "hi"}]) == "老路径"