zcbot/tests/test_media_common.py

"""tools/media_common.py 单测 —— 五个媒体工具同构析出后的共享原语。

全部纯函数/打桩,不碰网络与 DB(quota_gate 的 check_daily_quota、
ark_chat_with_retry 的 ArkClient 均 monkeypatch)。
"""
from __future__ import annotations

import re
import json
import unittest
import uuid
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest import mock

from core.ark_client import ArkError, ArkTimeoutError
from tools import media_common as mc


class StampedPathTests(unittest.TestCase):
    def test_creates_dir_and_names_with_ts_rand(self):
        with TemporaryDirectory() as td:
            d = Path(td) / "figures" / "sub"
            p = mc.stamped_path(d, ".png")
            self.assertTrue(d.is_dir())                      # 顺手建目录
            self.assertEqual(p.parent, d)
            self.assertRegex(p.name, r"^\d{8}-\d{6}-[0-9a-f]{6}\.png$")

    def test_write_meta_alongside(self):
        with TemporaryDirectory() as td:
            dest = Path(td) / "a.mp4"
            mc.write_meta(dest, {"prompt": "早安", "cost_cny": 0.22})
            meta = json.loads((Path(td) / "a.meta.json").read_text(encoding="utf-8"))
            self.assertEqual(meta["prompt"], "早安")
            self.assertEqual(meta["cost_cny"], 0.22)


class QuotaGateTests(unittest.TestCase):
    def test_no_limit_short_circuits(self):
        # limit<=0 不查 DB(不打桩也不该被调用)
        with mock.patch.object(mc, "check_daily_quota", side_effect=AssertionError("不该调")):
            self.assertIsNone(mc.quota_gate(uuid.uuid4(), kind="image", limit=0, what="图片生成", noun="张"))

    def test_over_and_under(self):
        uid = uuid.uuid4()
        with mock.patch.object(mc, "check_daily_quota", return_value=(30, True)):
            msg = mc.quota_gate(uid, kind="image", limit=30, what="图片生成", noun="张")
            self.assertEqual(msg, "[Error] 已达每日图片生成上限(30/30 张),次日 00:00 重置。")
        with mock.patch.object(mc, "check_daily_quota", return_value=(3, False)):
            self.assertIsNone(mc.quota_gate(uid, kind="video", limit=10, what="视频生成", noun="个"))


class RecordUsageSafeTests(unittest.TestCase):
    def test_returns_value(self):
        self.assertEqual(mc.record_usage_safe("t", lambda **kw: 42, a=1), 42)

    def test_swallow_exception_returns_none(self):
        def boom(**kw):
            raise RuntimeError("db down")
        self.assertIsNone(mc.record_usage_safe("t", boom))


class ExtractChatAnswerTests(unittest.TestCase):
    def test_str_content(self):
        resp = {"choices": [{"finish_reason": "stop", "message": {"content": "  你好  "}}]}
        self.assertEqual(mc.extract_chat_answer(resp), ("你好", False))

    def test_list_content_and_truncated(self):
        resp = {"choices": [{
            "finish_reason": "length",
            "message": {"content": [
                {"type": "text", "text": "第一段"},
                {"type": "image_url", "image_url": {}},
                {"type": "text", "text": "第二段"},
            ]},
        }]}
        self.assertEqual(mc.extract_chat_answer(resp), ("第一段\n第二段", True))

    def test_degenerate_shapes(self):
        self.assertEqual(mc.extract_chat_answer(None), ("", False))
        self.assertEqual(mc.extract_chat_answer({}), ("", False))
        self.assertEqual(mc.extract_chat_answer({"choices": []}), ("", False))
        self.assertEqual(
            mc.extract_chat_answer({"choices": [{"finish_reason": "length"}]}), ("", True)
        )


class FindFirstUrlTests(unittest.TestCase):
    def test_nested_image_url(self):
        resp = {"data": {"images": [{"url": "http://x/img.png"}]}}
        self.assertEqual(mc.find_first_url(resp, keys=("url", "image_url")), "http://x/img.png")

    def test_non_http_ignored(self):
        resp = {"url": "figures/local.png", "data": [{"image_url": "https://y/z.png"}]}
        self.assertEqual(mc.find_first_url(resp, keys=("url", "image_url")), "https://y/z.png")

    def test_accept_predicate_video(self):
        accept = lambda k, v: k == "video_url" or v.lower().endswith((".mp4", ".webm", ".mov"))
        # 裸 url 非视频扩展 → 拒;video_url → 收
        resp = {"url": "http://x/preview.png", "content": {"video_url": "http://x/v"}}
        self.assertEqual(
            mc.find_first_url(resp, keys=("video_url", "url"), accept=accept), "http://x/v"
        )
        # 只有裸视频扩展 url → 收
        resp2 = {"data": [{"url": "http://x/out.MP4"}]}
        self.assertEqual(
            mc.find_first_url(resp2, keys=("video_url", "url"), accept=accept), "http://x/out.MP4"
        )

    def test_not_found(self):
        self.assertEqual(mc.find_first_url({"a": [1, {"b": "no"}]}, keys=("url",)), "")


class _FakeClient:
    """打桩 ArkClient:按脚本依次返回/抛出。"""
    script: list = []

    def __init__(self, *a, **kw):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *a):
        return False

    def post_json(self, endpoint, body, timeout_s=None):
        act = _FakeClient.script.pop(0)
        if isinstance(act, Exception):
            raise act
        return act


class ArkChatWithRetryTests(unittest.TestCase):
    def _run(self, script, retries=1):
        _FakeClient.script = list(script)
        with mock.patch.object(mc, "ArkClient", _FakeClient), \
             mock.patch.object(mc.time, "sleep"):
            return mc.ark_chat_with_retry(
                object(), "/chat/completions", {}, timeout_s=1.0,
                retries=retries, tool_name="look_at_image",
            )

    def test_success_first_try(self):
        resp, err = self._run([{"ok": 1}])
        self.assertEqual(resp, {"ok": 1})
        self.assertEqual(err, "")

    def test_timeout_then_success(self):
        resp, err = self._run([ArkTimeoutError("t"), {"ok": 2}], retries=1)
        self.assertEqual(resp, {"ok": 2})
        self.assertEqual(err, "")

    def test_timeout_exhausted(self):
        resp, err = self._run([ArkTimeoutError("t1"), ArkTimeoutError("t2")], retries=1)
        self.assertIsNone(resp)
        self.assertIn("look_at_image API", err)
        self.assertIn("已重试 1 次仍超时", err)

    def test_business_error_no_retry(self):
        resp, err = self._run([ArkError("HTTP 400 bad"), {"ok": 3}], retries=3)
        self.assertIsNone(resp)
        self.assertIn("HTTP 400 bad", err)
        self.assertEqual(len(_FakeClient.script), 1)  # 第二个动作没被消费 = 未重试


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