"""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_in_hidden_directory(self): with TemporaryDirectory() as td: dest = Path(td) / "a.mp4" mc.write_meta(dest, {"prompt": "早安", "cost_cny": 0.22}) meta_path = Path(td) / ".meta" / "a.mp4.json" meta = json.loads(meta_path.read_text(encoding="utf-8")) self.assertEqual(meta["prompt"], "早安") self.assertEqual(meta["cost_cny"], 0.22) self.assertFalse((Path(td) / "a.meta.json").exists()) 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()