340 lines
12 KiB
Python
340 lines
12 KiB
Python
"""
|
|
私聊附件询价：发卡、H5/图并入、口播复用文字流程。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if ROOT not in sys.path:
|
|
sys.path.insert(0, ROOT)
|
|
|
|
os.environ.setdefault("YTD_ENV", "test")
|
|
os.environ.setdefault("REDIS_BACKEND", "memory")
|
|
os.environ.setdefault("CHECKPOINT_BACKEND", "memory")
|
|
os.environ.setdefault("MESSAGE_STORE_BACKEND", "memory")
|
|
os.environ.setdefault("LLM_DATA_USAGE_CONFIRMED", "false")
|
|
os.environ.setdefault("LLM_ALLOW_NETWORK", "false")
|
|
os.environ.setdefault("LEDGER_BACKEND", "memory")
|
|
|
|
from agent.channel.queue import MemoryMessageStore
|
|
from agent.channel.wecom.models import InboundMessage
|
|
from agent.handlers import attachment_inquiry as attach_mod
|
|
from agent.handlers.attachment_inquiry import (
|
|
accept_h5_uploads,
|
|
finish_queued_attachment,
|
|
handle_attachment_inquiry,
|
|
join_attachment_image,
|
|
)
|
|
from agent.handlers.text_inquiry import handle_text_inquiry
|
|
from agent.policy import inquiry_copy as copy
|
|
from agent.policy.air_text_flow import reset_air_text_flow_for_test
|
|
from agent.policy.attachment_wave import get_attach_wave_book, reset_attach_wave_book_for_test
|
|
from agent.policy.image_wave import reset_image_wave_book_for_test
|
|
from agent.policy.land_text_flow import reset_land_text_flow_for_test
|
|
from agent.policy.sea_text_flow import get_sea_text_flow, reset_sea_text_flow_for_test
|
|
from agent.routing.decision import RouteDecision
|
|
from agent.routing.dispatch import dispatch_inbound
|
|
from tests.test_sea_text_inquiry import COMPLETE_SEA
|
|
|
|
SEA_ORAL = """海运
|
|
起运港：上海
|
|
目的港：洛杉矶
|
|
品名：普货
|
|
货量：20吨
|
|
整柜或拼柜：整柜
|
|
箱型箱量：1x40HQ
|
|
"""
|
|
|
|
|
|
def _extract(text: str):
|
|
def _fn(_name: str, _raw: bytes):
|
|
body = (text or "").strip()
|
|
return {"ok": True, "text": body, "empty": not bool(body), "error": ""}
|
|
|
|
return _fn
|
|
|
|
|
|
def _vision(text: str, *, ok: bool = True, error: str = ""):
|
|
def _fn(**_kwargs):
|
|
body = (text or "").strip()
|
|
return {"ok": ok, "text": body, "empty": not bool(body), "error": error}
|
|
|
|
return _fn
|
|
|
|
|
|
def _img(sender: str, mid: str) -> InboundMessage:
|
|
return InboundMessage(
|
|
sender_id=sender,
|
|
message_id=mid,
|
|
content="",
|
|
msg_type="image",
|
|
media={"media_id": mid, "filename": f"{mid}.jpg"},
|
|
)
|
|
|
|
|
|
class AttachmentInquiryTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
reset_sea_text_flow_for_test()
|
|
reset_land_text_flow_for_test()
|
|
reset_air_text_flow_for_test()
|
|
reset_attach_wave_book_for_test()
|
|
reset_image_wave_book_for_test()
|
|
self.store = MemoryMessageStore()
|
|
|
|
def _bodies(self) -> list[str]:
|
|
return [i.content for i in self.store._outbox.values()]
|
|
|
|
def _extras(self) -> list[dict]:
|
|
return [i.payload for i in self.store._outbox.values()]
|
|
|
|
def test_intent_sends_upload_card_not_ticket(self) -> None:
|
|
phase = handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="帮我传个Excel询价，海运上海到纽约"),
|
|
store=self.store,
|
|
)
|
|
self.assertEqual(phase, "wait_upload")
|
|
self.assertIsNone(get_sea_text_flow().session_of("s1"))
|
|
extras = self._extras()
|
|
self.assertTrue(extras)
|
|
card = extras[0].get("template_card") or {}
|
|
self.assertEqual(card.get("card_type"), "text_notice")
|
|
self.assertEqual(card.get("main_title", {}).get("title"), "上传询价附件")
|
|
self.assertIn("上传询价附件", (card.get("jump_list") or [{}])[0].get("title") or "")
|
|
url = (card.get("card_action") or {}).get("url") or ""
|
|
self.assertIn("/inquiry-form/upload/", url)
|
|
|
|
def test_xlsx_complete_creates_sea_ticket(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
phase, hint = accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("sea.xlsx", b"fake")],
|
|
store=self.store,
|
|
extract_fn=_extract(SEA_ORAL),
|
|
attach_text_fn=lambda **k: {
|
|
"ok": True,
|
|
"text": str((k.get("payload") or {}).get("raw_text") or ""),
|
|
"empty": False,
|
|
"error": "",
|
|
},
|
|
put_fn=lambda **kw: {"ok": True, "key": kw.get("key") or "k"},
|
|
)
|
|
self.assertEqual(hint, copy.ATTACH_H5_OK)
|
|
self.assertIn(copy.ATTACH_RECOGNIZING, self._bodies())
|
|
sess = get_sea_text_flow().session_of("s1")
|
|
self.assertIsNotNone(sess)
|
|
self.assertTrue((sess.work_order_no or "").strip())
|
|
ticket = get_sea_text_flow().ledger.get_ticket(work_order_no=sess.work_order_no)
|
|
self.assertTrue(ticket.attachments)
|
|
|
|
def test_empty_file_asks_mode(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("empty.xlsx", b"x")],
|
|
store=self.store,
|
|
extract_fn=_extract(""),
|
|
attach_text_fn=lambda **_k: {"ok": True, "text": "", "empty": True, "error": ""},
|
|
)
|
|
bodies = self._bodies()
|
|
self.assertTrue(any("没看清" in b for b in bodies))
|
|
self.assertTrue(any("海运" in b and "空运" in b for b in bodies))
|
|
|
|
def test_image_during_wait_joins_same_ticket(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
join_attachment_image(
|
|
_img("s1", "mid-1"),
|
|
store=self.store,
|
|
vision_fn=_vision("截图"),
|
|
download_fn=lambda _m: (b"x", ""),
|
|
process_now=False,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("sea.xlsx", b"x")],
|
|
store=self.store,
|
|
extract_fn=_extract(SEA_ORAL),
|
|
attach_text_fn=lambda **k: {
|
|
"ok": True,
|
|
"text": str((k.get("payload") or {}).get("raw_text") or ""),
|
|
"empty": False,
|
|
},
|
|
process_now=True,
|
|
)
|
|
from agent.handlers.attachment_inquiry import finish_attach_wave
|
|
from agent.policy.attachment_wave import get_attach_wave_book
|
|
|
|
book = get_attach_wave_book()
|
|
book.complete_material("s1", 0, text="截图")
|
|
finish_attach_wave("s1", store=self.store)
|
|
acks = [c for c in self._bodies() if c in {copy.IMAGE_RECOGNIZING, copy.ATTACH_RECOGNIZING}]
|
|
self.assertEqual(len(acks), 1)
|
|
self.assertEqual(acks[0], copy.IMAGE_RECOGNIZING)
|
|
sess = get_sea_text_flow().session_of("s1")
|
|
self.assertTrue(sess and sess.work_order_no)
|
|
|
|
def test_mid_clarify_then_attach_is_new_ticket(self) -> None:
|
|
handle_text_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="old", content="海运出号"),
|
|
store=self.store,
|
|
injected_facts=COMPLETE_SEA,
|
|
injected_mode="SEA",
|
|
)
|
|
wo_old = get_sea_text_flow().session_of("s1").work_order_no
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m2", content="我有个excel要询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("sea.xlsx", b"x")],
|
|
store=self.store,
|
|
extract_fn=_extract(SEA_ORAL),
|
|
attach_text_fn=lambda **k: {
|
|
"ok": True,
|
|
"text": str((k.get("payload") or {}).get("raw_text") or ""),
|
|
"empty": False,
|
|
},
|
|
)
|
|
wo_new = get_sea_text_flow().session_of("s1").work_order_no
|
|
self.assertNotEqual(wo_old, wo_new)
|
|
self.assertIsNotNone(get_sea_text_flow().ledger.get_ticket(work_order_no=wo_old))
|
|
|
|
def test_reject_png_on_h5(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
phase, hint = accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("shot.png", b"\x89PNG")],
|
|
store=self.store,
|
|
)
|
|
self.assertEqual(phase, "reject_type")
|
|
self.assertIn("Excel", hint)
|
|
self.assertIsNone(get_sea_text_flow().session_of("s1"))
|
|
|
|
def test_dispatch_image_joins_wait_not_new_image_inquiry(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
name = dispatch_inbound(
|
|
_img("s1", "mid-x"),
|
|
RouteDecision(intent="image_inquiry"),
|
|
)
|
|
self.assertEqual(name, "attachment_inquiry")
|
|
|
|
def test_follow_text_not_create_ticket_before_upload(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
name = dispatch_inbound(
|
|
InboundMessage(sender_id="s1", message_id="m2", content="海运上海到纽约 2个40尺"),
|
|
RouteDecision(intent="ordinary_text_inquiry"),
|
|
)
|
|
self.assertEqual(name, "attach_wave_follow_text")
|
|
self.assertIsNone(get_sea_text_flow().session_of("s1"))
|
|
|
|
def test_after_business_old_h5_rejected(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("sea.xlsx", b"x")],
|
|
store=self.store,
|
|
extract_fn=_extract(SEA_ORAL),
|
|
attach_text_fn=lambda **k: {
|
|
"ok": True,
|
|
"text": str((k.get("payload") or {}).get("raw_text") or ""),
|
|
"empty": False,
|
|
},
|
|
)
|
|
phase, _hint = accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("again.xlsx", b"x")],
|
|
store=self.store,
|
|
extract_fn=_extract(SEA_ORAL),
|
|
)
|
|
self.assertEqual(phase, "expired")
|
|
|
|
def test_worker_reads_bytes_after_http_stash_cleared(self) -> None:
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("sea.xlsx", b"fake-xlsx-bytes")],
|
|
store=self.store,
|
|
process_now=False,
|
|
)
|
|
attach_mod._STASH.clear()
|
|
wave = get_attach_wave_book().active("s1")
|
|
self.assertIsNotNone(wave)
|
|
key = f"inquiry-attach/s1/{wave.wave_id}/0/sea.xlsx"
|
|
finish_queued_attachment(
|
|
{
|
|
"sender_id": "s1",
|
|
"index": 0,
|
|
"role": "file",
|
|
"filename": "sea.xlsx",
|
|
"object_key": key,
|
|
"wave_id": wave.wave_id,
|
|
"__extract": _extract(SEA_ORAL),
|
|
"__attach_text": lambda **k: {
|
|
"ok": True,
|
|
"text": str((k.get("payload") or {}).get("raw_text") or ""),
|
|
"empty": False,
|
|
"error": "",
|
|
},
|
|
}
|
|
)
|
|
self.assertIsNotNone(get_sea_text_flow().session_of("s1"))
|
|
|
|
def test_second_unreadable_wave_still_replies(self) -> None:
|
|
def _empty(_name: str, _raw: bytes):
|
|
return {"ok": True, "text": "", "empty": True, "error": ""}
|
|
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m1", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("blank.xlsx", b"xx")],
|
|
store=self.store,
|
|
extract_fn=_empty,
|
|
)
|
|
handle_attachment_inquiry(
|
|
InboundMessage(sender_id="s1", message_id="m2", content="附件询价"),
|
|
store=self.store,
|
|
)
|
|
accept_h5_uploads(
|
|
sender_id="s1",
|
|
files=[("blank2.xlsx", b"yy")],
|
|
store=self.store,
|
|
extract_fn=_empty,
|
|
)
|
|
fails = [b for b in self._bodies() if "没看清" in b or "没认出来" in b]
|
|
self.assertGreaterEqual(len(fails), 2, fails)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|