inquiry_robot/inquiry-agent/tests/test_attachment_wave.py

"""
私聊附件等材料波次：发卡等待、发图/H5 并入、30 分钟时效。
"""

from __future__ import annotations

import json
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)

from agent.policy.attachment_wave import (
    WAIT_TTL_SECONDS,
    AttachWaveBook,
    reset_attach_wave_book_for_test,
)


class AttachWaveTests(unittest.TestCase):
    def setUp(self) -> None:
        self.book = reset_attach_wave_book_for_test()

    def test_wait_ttl_is_30_minutes(self) -> None:
        self.assertEqual(WAIT_TTL_SECONDS, 1800)

    def test_open_wait_is_active_without_files(self) -> None:
        wave = self.book.open_wait("s1", precede_text="海运上海到纽约", token="tok1", now=100.0)
        self.assertTrue(wave.wave_id)
        self.assertEqual(wave.precede_text, "海运上海到纽约")
        self.assertFalse(wave.started)
        self.assertIsNotNone(self.book.active("s1", now=100.0))

    def test_wait_expires_after_30_minutes_if_not_started(self) -> None:
        self.book.open_wait("s1", precede_text="海运", token="tok1", now=100.0)
        self.assertIsNone(self.book.active("s1", now=100.0 + 1800.1))

    def test_follow_text_while_waiting(self) -> None:
        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.assertTrue(self.book.append_follow_text("s1", "只要整柜"))
        self.assertFalse(self.book.append_follow_text("s1", ""))

    def test_join_image_starts_recognition(self) -> None:
        self.book.open_wait("s1", precede_text="海运", token="tok1", now=1.0)
        wave, how = self.book.join_image("s1", "mid-a", filename="a.jpg", now=2.0)
        self.assertEqual(how, "joined")
        self.assertTrue(wave.started)
        self.assertEqual(self.book.pending(wave), 1)

    def test_add_files_keeps_upload_order(self) -> None:
        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        wave = self.book.add_files("s1", ["sea.xlsx", "air.pdf"], now=2.0)
        self.assertTrue(wave.started)
        self.assertEqual([m.filename for m in wave.materials], ["sea.xlsx", "air.pdf"])
        self.assertEqual(self.book.pending(wave), 2)

    def test_compose_oral_uses_material_order(self) -> None:
        self.book.open_wait("s1", precede_text="开口字", token="tok1", now=1.0)
        self.book.add_files("s1", ["a.xlsx", "b.pdf"], now=2.0)
        self.book.append_follow_text("s1", "只要40尺")
        wave = self.book.active("s1", now=2.0)
        self.assertIsNotNone(wave)
        self.book.complete_material("s1", 1, text="空运 上海到纽约")
        self.book.complete_material("s1", 0, text="海运 宁波到汉堡")
        oral = self.book.compose_oral("s1")
        self.assertEqual(oral, "开口字\n海运 宁波到汉堡\n空运 上海到纽约\n只要40尺")

    def test_empty_land_quote_sheet_still_makes_multi_with_sea(self) -> None:
        """陆运报价单被整理成空时，仍要按文件名留陆运段，不能只剩海运。"""
        from agent.policy.multi_segments import classify_multi_text

        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.book.add_files(
            "s1",
            ["WO202608240017_陆运报价单(1).xlsx", "手工报价海运.xlsx"],
            now=2.0,
        )
        self.book.complete_material("s1", 0, text="")
        self.book.complete_material(
            "s1",
            1,
            text="运输方式：海运\n起运港：纽约港\n目的港：巴生西港",
        )
        oral = self.book.compose_oral("s1")
        self.assertIn("陆运", oral)
        self.assertIn("海运", oral)
        self.assertIn("纽约港", oral)
        self.assertEqual(classify_multi_text(oral), "multi")

    def test_one_sea_quote_filename_is_single(self) -> None:
        """只传一份海运报价单：文件名带海运，表里再写运输方式，不能报同一方式两段。"""
        from agent.policy.multi_segments import classify_multi_text

        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.book.add_files("s1", ["WO202608250012_海运报价单.xlsx"], now=2.0)
        self.book.complete_material(
            "s1",
            0,
            text="运输方式：海运\n起运港：纽约港\n目的港：巴生西港\n海运询价",
        )
        oral = self.book.compose_oral("s1")
        self.assertIn("段1：海运", oral)
        self.assertIn("运输方式：海运", oral)
        self.assertEqual(classify_multi_text(oral), "single")

    def test_ltl_filename_counts_as_land_segment(self) -> None:
        """零担询价表文件名没有「陆运」二字，也要和海运表切成两段。"""
        from agent.policy.multi_segments import classify_multi_text

        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.book.add_files(
            "s1",
            ["东莞-上海  整车及零担 询价.xlsx", "手工海运报价单.xlsx"],
            now=2.0,
        )
        self.book.complete_material("s1", 0, text="始发站：东莞\n目的站：上海")
        self.book.complete_material("s1", 1, text="起运港：宁波港")
        oral = self.book.compose_oral("s1")
        self.assertIn("陆运", oral)
        self.assertIn("海运", oral)
        self.assertEqual(classify_multi_text(oral), "multi")

    def test_after_business_done_not_active(self) -> None:
        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.book.add_files("s1", ["a.xlsx"], now=2.0)
        self.book.complete_material("s1", 0, text="海运")
        self.book.mark_business_done("s1")
        self.assertIsNone(self.book.active("s1", now=3.0))
        self.assertFalse(self.book.append_follow_text("s1", "补一句"))

    def test_image_after_business_rejected(self) -> None:
        self.book.open_wait("s1", precede_text="", token="tok1", now=1.0)
        self.book.mark_business_done("s1")
        wave, how = self.book.join_image("s1", "mid-b", filename="b.jpg", now=2.0)
        self.assertEqual(how, "inactive")
        self.assertIsNone(wave)

    def test_http_and_worker_share_kv(self) -> None:
        kv = _MemoryWaveKV()
        http = reset_attach_wave_book_for_test(kv=kv)
        worker = AttachWaveBook(kv=kv)
        http.open_wait("s1", precede_text="海运", token="tok", now=1.0)
        self.assertIsNotNone(worker.active("s1", now=1.0))
        worker.add_files("s1", ["a.xlsx"], now=2.0)
        worker.complete_material("s1", 0, text="宁波到汉堡")
        worker.mark_business_done("s1")
        self.assertIsNone(http.active("s1", now=3.0))


class _MemoryWaveKV:
    """单测假 Redis：两份账本共用同一字典。"""

    def __init__(self) -> None:
        self._d: dict[str, dict] = {}

    def get_json(self, kind: str, sender_id: str):
        row = self._d.get(f"{kind}:{sender_id}")
        return json.loads(json.dumps(row)) if row is not None else None

    def set_json(self, kind: str, sender_id: str, data: dict) -> None:
        self._d[f"{kind}:{sender_id}"] = json.loads(json.dumps(data))

    def delete(self, kind: str, sender_id: str) -> None:
        self._d.pop(f"{kind}:{sender_id}", None)


if __name__ == "__main__":
    unittest.main()