Files
cmshoppe/tests/test_ai.py
T

2239 lines
94 KiB
Python
Raw Normal View History

2026-06-27 15:47:21 +08:00
import base64
import io
import json
import os
2026-07-04 15:13:15 +08:00
import socket
2026-06-27 15:47:21 +08:00
import sys
2026-07-07 21:00:43 +08:00
import threading
2026-06-27 15:47:21 +08:00
import unittest
from concurrent.futures import CancelledError
2026-06-27 16:26:08 +08:00
from types import SimpleNamespace
2026-06-27 15:47:21 +08:00
from unittest import mock
sys.path.insert(0, os.path.dirname(__file__))
from _helpers import TempDirMixin
2026-06-27 16:26:08 +08:00
from app import ai, appconfig, db
2026-06-27 15:47:21 +08:00
class _Response:
def __init__(self, payload):
self.payload = payload
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
def read(self, size=-1):
return json.dumps(self.payload).encode("utf-8")
2026-07-04 15:13:15 +08:00
class _RequestsResponse:
def __init__(self, payload=None, status_code=200, content=b"", headers=None):
self.payload = payload if payload is not None else {}
self.status_code = status_code
self.content = content
self.headers = headers or {}
self.text = json.dumps(self.payload, ensure_ascii=False)
self.closed = False
2026-07-04 15:13:15 +08:00
def json(self):
return self.payload
def iter_content(self, chunk_size=65536):
if self.content:
yield self.content
def close(self):
self.closed = True
2026-06-27 15:47:21 +08:00
class AITests(TempDirMixin, unittest.TestCase):
def _write_models(self, path, text=None, image=None):
text = text or {
"name": "Text",
"category": "text",
"enabled": True,
"url": "https://example.invalid/v1/chat/completions",
"model": "text-model",
"api_key": "sk-text-secret",
"api_type": "chat",
"connect_timeout_seconds": 1,
"timeout_seconds": 1,
"extra_body": {"temperature": 0},
}
image = image or {
"name": "Image",
"category": "image",
"enabled": True,
"url": "https://example.invalid/v1/chat/completions",
"model": "image-model",
"api_key": "sk-image-secret",
"api_type": "auto",
"connect_timeout_seconds": 1,
"timeout_seconds": 1,
"extra_body": {},
}
appconfig.save_ai_models_config({"models": [text, image]}, path=path)
def _config(self):
cfg = appconfig.default_config()
2026-07-04 17:59:23 +08:00
cfg["ai"]["backend"] = "direct"
2026-06-27 15:47:21 +08:00
cfg["ai"]["default_text_model"] = "Text"
cfg["ai"]["default_image_model"] = "Image"
cfg["ai"]["retry"] = 1
cfg["ai"]["resolution"] = "512"
cfg["ai"]["jpg_quality"] = 80
return cfg
2026-07-04 15:13:15 +08:00
def _cmhub_config(self, temp_dir):
cfg = self._config()
cfg["ai"]["backend"] = "cmhub"
cfg["ai"]["resolution"] = "512"
cfg["ai"]["cmhub"] = {
"base_url": "https://cmhub.example.com",
"title_alias": "title-standard",
"image_alias": "image-hd",
"vision_alias": "vision-standard",
2026-07-04 15:13:15 +08:00
"connect_timeout": 3,
"download_with_curl": "false",
2026-07-04 15:13:15 +08:00
"check_balance_before_batch": False,
}
key_path = os.path.join(temp_dir, "cmhub.json")
appconfig.save_cmhub_config({"api_key": "sk-cmhub-secret"}, path=key_path)
return cfg, key_path
2026-06-27 16:26:08 +08:00
def _collected_tasks(self, temp_dir, cfg, titles=None):
titles = titles or ["旧标题A", "旧标题B"]
db.init_db(cfg["db_path"])
batch_id = db.create_batch([os.path.join(temp_dir, "input.xlsx")], path=cfg["db_path"])
db.insert_tasks(
batch_id,
[
{
"source_file_abs": os.path.join(temp_dir, "input.xlsx"),
"source_sheet": "商品",
"source_row": index + 2,
"account_name": "Excel主店",
"alias": "alias-a",
"item_id": "5110063951%s" % index,
}
for index in range(len(titles))
],
path=cfg["db_path"],
)
tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
for task, title in zip(tasks, titles):
db.set_collected(
task.id,
title,
os.path.join(temp_dir, "%s_old.jpg" % task.item_id),
path=cfg["db_path"],
)
return batch_id, db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
def test_cover_title_context_and_cover_only_generation_needs(self):
task = SimpleNamespace(
stage="collected",
status="success",
old_title="旧标题",
new_title=None,
new_cover_path=None,
apply_attempts=0,
)
self.assertEqual("旧标题", ai.cover_title_context(task))
self.assertEqual(
{"title": False, "cover": True},
ai.generation_needs(task, generate_mode="cover"),
)
task.new_title = "新标题"
self.assertEqual("新标题", ai.cover_title_context(task))
task.new_title = None
task.old_title = ""
self.assertEqual("", ai.cover_title_context(task))
self.assertFalse(ai.is_generatable_task(task, generate_mode="cover"))
self.assertEqual(
{"title": True, "cover": True},
ai.generation_needs(task, generate_mode="title_cover"),
)
2026-07-07 21:00:43 +08:00
def _write_old_cover_files(self, tasks):
from PIL import Image
for index, task in enumerate(tasks):
os.makedirs(os.path.dirname(task.old_cover_path), exist_ok=True)
Image.new(
"RGB",
(16, 16),
((index * 35) % 255, 80, 120),
).save(task.old_cover_path, "JPEG")
def _png_bytes(self, color=(200, 120, 80)):
from PIL import Image
generated = io.BytesIO()
Image.new("RGB", (8, 8), color).save(generated, "PNG")
return generated.getvalue()
2026-06-27 15:47:21 +08:00
def test_gen_title_uses_configured_model_and_retries(self):
with self.make_temp_dir() as temp_dir:
models_path = os.path.join(temp_dir, "ai_models.json")
self._write_models(models_path)
calls = []
steps = []
2026-06-27 15:47:21 +08:00
def fake_urlopen(request, timeout=None):
calls.append((request, timeout))
if len(calls) == 1:
raise ai.urllib.error.URLError("temporary")
return _Response({"choices": [{"message": {"content": " 新标题 "}}]})
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
title = ai.gen_title(
"优化标题",
"旧标题",
config=self._config(),
models_path=models_path,
on_step=steps.append,
2026-06-27 15:47:21 +08:00
)
self.assertEqual("新标题", title)
self.assertEqual(2, len(calls))
retry_events = [step for step in steps if isinstance(step, dict)]
self.assertEqual(1, len(retry_events))
self.assertEqual("title_request", retry_events[0]["step"])
self.assertEqual("retry", retry_events[0]["result"])
self.assertEqual(1, retry_events[0]["attempt"])
self.assertEqual(2, retry_events[0]["attempts"])
self.assertNotIn("sk-text-secret", retry_events[0]["detail"])
2026-06-27 15:47:21 +08:00
body = json.loads(calls[-1][0].data.decode("utf-8"))
self.assertEqual("text-model", body["model"])
self.assertEqual(0, body["temperature"])
self.assertEqual(1, len(body["messages"]))
self.assertIn("优化标题", body["messages"][0]["content"])
self.assertIn("旧标题:\n旧标题", body["messages"][0]["content"])
self.assertIn("请只返回新标题,不要解释。", body["messages"][0]["content"])
self.assertNotIn("sk-text-secret", body["messages"][0]["content"])
self.assert_removed(temp_dir)
def test_gen_title_renders_old_title_placeholder_without_duplicate_append(self):
with self.make_temp_dir() as temp_dir:
models_path = os.path.join(temp_dir, "ai_models.json")
self._write_models(models_path)
calls = []
def fake_urlopen(request, timeout=None):
calls.append(request)
return _Response({"choices": [{"message": {"content": "新标题"}}]})
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
title = ai.gen_title(
"请基于{旧标题}重写标题,保留 {商品id}",
"原始标题A",
config=self._config(),
models_path=models_path,
)
self.assertEqual("新标题", title)
body = json.loads(calls[-1].data.decode("utf-8"))
content = body["messages"][0]["content"]
self.assertIn("请基于原始标题A重写标题", content)
self.assertEqual(1, content.count("原始标题A"))
self.assertNotIn("旧标题:", content)
self.assertIn("{商品id}", content)
self.assertIn("请只返回新标题,不要解释。", content)
2026-06-27 15:47:21 +08:00
self.assert_removed(temp_dir)
def test_gen_title_accepts_openai_compatible_base_url(self):
with self.make_temp_dir() as temp_dir:
models_path = os.path.join(temp_dir, "ai_models.json")
self._write_models(
models_path,
text={
"name": "Text",
"category": "text",
"enabled": True,
"url": "https://example.invalid/v1",
"model": "text-model",
"api_key": "sk-text-secret",
"api_type": "chat",
"connect_timeout_seconds": 1,
"timeout_seconds": 1,
"extra_body": {},
},
)
calls = []
steps = []
def fake_urlopen(request, timeout=None):
calls.append(request)
return _Response({"choices": [{"message": {"content": "新标题"}}]})
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
title = ai.gen_title(
"优化标题",
"旧标题",
config=self._config(),
models_path=models_path,
on_step=steps.append,
)
self.assertEqual("新标题", title)
self.assertEqual(
"https://example.invalid/v1/chat/completions",
calls[0].full_url,
)
self.assert_removed(temp_dir)
2026-06-27 15:47:21 +08:00
def test_missing_model_fields_raise_clear_error_without_secret(self):
with self.make_temp_dir() as temp_dir:
models_path = os.path.join(temp_dir, "ai_models.json")
self._write_models(
models_path,
text={
"name": "Text",
"category": "text",
"enabled": True,
"url": "",
"model": "text-model",
"api_key": "sk-text-secret",
"api_type": "chat",
"connect_timeout_seconds": 1,
"timeout_seconds": 1,
"extra_body": {},
},
)
with self.assertRaises(ai.AIError) as raised:
ai.gen_title("prompt", "old", config=self._config(), models_path=models_path)
message = str(raised.exception)
self.assertIn("缺少字段: url", message)
self.assertNotIn("sk-text-secret", message)
self.assert_removed(temp_dir)
def test_gen_cover_saves_jpeg_with_resolution_and_quality(self):
try:
from PIL import Image
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
models_path = os.path.join(temp_dir, "ai_models.json")
self._write_models(models_path)
old_cover = os.path.join(temp_dir, "old.jpg")
output = os.path.join(temp_dir, "new.jpg")
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
generated = io.BytesIO()
Image.new("RGB", (8, 8), (200, 120, 80)).save(generated, "PNG")
b64_image = base64.b64encode(generated.getvalue()).decode("ascii")
def fake_urlopen(request, timeout=None):
body = json.loads(request.data.decode("utf-8"))
self.assertEqual("image-model", body["model"])
self.assertIn("目标分辨率:512", body["messages"][0]["content"][0]["text"])
return _Response({"data": [{"b64_json": b64_image}]})
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
result = ai.gen_cover(
"生成封面",
old_cover,
output,
config=self._config(),
models_path=models_path,
)
self.assertEqual(os.path.abspath(output), result)
with Image.open(output) as saved:
self.assertEqual((512, 512), saved.size)
self.assertEqual("JPEG", saved.format)
self.assert_removed(temp_dir)
2026-07-04 15:13:15 +08:00
def test_cmhub_gen_title_uses_alias_and_emits_metadata(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
calls = []
events = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse(
{
"titles": [" 新标题 "],
"alias": "title-standard",
"model_used": "provider-title-model",
"points_cost": 1,
"points_balance": 99,
"call_id": "call-title-1",
}
)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
2026-07-04 15:13:15 +08:00
title = ai.gen_title(
"优化标题",
"旧标题",
config=cfg,
cmhub_config_path=key_path,
on_event=events.append,
)
self.assertEqual("新标题", title)
self.assertEqual(1, len(calls))
method, url, kwargs = calls[0]
self.assertEqual("POST", method)
self.assertEqual("https://cmhub.example.com/api/v1/generate/title", url)
2026-07-06 22:51:28 +08:00
self.assertEqual((3, 600), kwargs["timeout"])
2026-07-04 15:13:15 +08:00
self.assertEqual("Bearer sk-cmhub-secret", kwargs["headers"]["Authorization"])
payload = kwargs["json"]
self.assertEqual("title-standard", payload["model"])
self.assertEqual("512", payload["resolution"])
self.assertIn("优化标题", payload["prompt"])
self.assertIn("旧标题", payload["prompt"])
self.assertTrue(events)
self.assertEqual("meta", events[0]["result"])
self.assertEqual(99, events[0]["metadata"]["points_balance"])
self.assert_removed(temp_dir)
def test_cmhub_gen_title_renders_old_title_placeholder_once(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse({"titles": ["新标题"]})
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
title = ai.gen_title(
"请参考 {旧标题} 输出更短标题,店铺变量 {店铺}",
"原始标题B",
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("新标题", title)
prompt = calls[0][2]["json"]["prompt"]
self.assertIn("请参考 原始标题B 输出更短标题", prompt)
self.assertEqual(1, prompt.count("原始标题B"))
self.assertNotIn("旧标题:", prompt)
self.assertIn("{店铺}", prompt)
self.assertIn("请只返回新标题,不要解释。", prompt)
self.assert_removed(temp_dir)
2026-07-04 15:13:15 +08:00
def test_cmhub_missing_config_raises_clear_error(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["ai"]["backend"] = "cmhub"
cfg["ai"]["cmhub"] = {
"base_url": "",
"title_alias": "",
"image_alias": "",
"connect_timeout": 3,
}
with self.assertRaises(ai.CMHubError) as raised:
ai.gen_title(
"prompt",
"old",
config=cfg,
cmhub_config_path=os.path.join(temp_dir, "missing.json"),
)
self.assertEqual("cmhub_not_configured", raised.exception.code)
self.assertIn("请去设置配置 cmhub", str(raised.exception))
2026-07-04 15:13:15 +08:00
self.assert_removed(temp_dir)
def test_cmhub_analyze_product_images_uses_vision_alias_and_safe_metadata(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
first = os.path.join(temp_dir, "first.png")
second = os.path.join(temp_dir, "second.jpg")
with open(first, "wb") as fh:
fh.write(b"first-image")
with open(second, "wb") as fh:
fh.write(b"second-image")
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse(
{
"text": "浅绿色连帽上衣,突出宽松版型与日常穿搭场景。",
"alias": "vision-standard",
"model_used": "vision-provider",
"points_cost": 1,
"points_balance": 231,
"call_id": "vision-call-1",
}
)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
result = ai.analyze_product_images(
"补充卖点:避免夸大",
"输出语言:繁体中文",
[first, second],
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("浅绿色连帽上衣,突出宽松版型与日常穿搭场景。", result["text"])
self.assertEqual(2, result["image_count"])
self.assertEqual(1, result["metadata"]["points_cost"])
self.assertEqual(231, result["metadata"]["points_balance"])
self.assertNotIn("image_base64", result)
self.assertEqual("POST", calls[0][0])
self.assertEqual(1, len(calls))
self.assertEqual(
"https://cmhub.example.com/api/v1/analyze/images",
calls[0][1],
)
payload = calls[0][2]["json"]
self.assertEqual("vision-standard", payload["model"])
self.assertEqual(2, len(payload["images"]))
self.assertTrue(payload["images"][0]["image_base64"].startswith("data:image/png;base64,"))
self.assertTrue(payload["images"][1]["image_base64"].startswith("data:image/jpeg;base64,"))
self.assertEqual({"temperature": 0.2}, payload["parameters"])
prompt = payload["prompt"]
self.assertIn("同一个商品项目", prompt)
self.assertIn("一组证据", prompt)
self.assertIn("禁止输出“图1/图2/第N张”", prompt)
self.assertIn("商品概述、可确认卖点、适用人群与场景、套图画面要求", prompt)
self.assertIn("待确认或避免编造的信息", prompt)
self.assertIn("可见差异", prompt)
self.assertEqual((3, ai.CMHUB_VISION_READ_TIMEOUT_SECONDS), calls[0][2]["timeout"])
self.assert_removed(temp_dir)
def test_cmhub_analyze_product_images_requires_vision_alias_and_local_limits(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
source = os.path.join(temp_dir, "source.jpg")
with open(source, "wb") as fh:
fh.write(b"source")
cfg["ai"]["cmhub"]["vision_alias"] = ""
with self.assertRaises(ai.CMHubError) as raised:
ai.analyze_product_images(
"提示",
"输出语言:繁体中文",
[source],
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("cmhub_not_configured", raised.exception.code)
self.assertIn("图片理解别名", str(raised.exception))
cfg["ai"]["cmhub"]["vision_alias"] = "vision-standard"
paths = []
for index in range(ai.CMHUB_VISION_MAX_IMAGES + 1):
path = os.path.join(temp_dir, "source-%d.jpg" % index)
with open(path, "wb") as fh:
fh.write(b"image")
paths.append(path)
with mock.patch.object(ai._cmhub_session(), "request") as request:
with self.assertRaises(ai.AIError) as too_many:
ai.analyze_product_images(
"提示",
"输出语言:繁体中文",
paths,
config=cfg,
cmhub_config_path=key_path,
)
self.assertIn("最多支持", str(too_many.exception))
request.assert_not_called()
oversized = os.path.join(temp_dir, "oversized.jpg")
with open(oversized, "wb") as fh:
fh.truncate(ai.CMHUB_VISION_MAX_IMAGE_BYTES + 1)
with mock.patch.object(ai._cmhub_session(), "request") as request:
with self.assertRaises(ai.AIError) as too_large:
ai.analyze_product_images(
"提示",
"输出语言:繁体中文",
[oversized],
config=cfg,
cmhub_config_path=key_path,
)
self.assertIn("超过10MiB", str(too_large.exception))
request.assert_not_called()
self.assert_removed(temp_dir)
def test_cmhub_analyze_product_images_read_timeout_is_not_retried(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
source = os.path.join(temp_dir, "source.jpg")
with open(source, "wb") as fh:
fh.write(b"source")
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
raise ai.requests.exceptions.ReadTimeout("slow")
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
with self.assertRaises(ai.CMHubError) as raised:
ai.analyze_product_images(
"提示",
"输出语言:繁体中文",
[source],
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("read_timeout", raised.exception.code)
self.assertIn("结果未确认", str(raised.exception))
self.assertEqual(1, len(calls))
self.assertEqual((3, ai.CMHUB_VISION_READ_TIMEOUT_SECONDS), calls[0][2]["timeout"])
self.assert_removed(temp_dir)
2026-07-04 15:13:15 +08:00
def test_cmhub_gen_cover_downloads_image_url_safely(self):
try:
from PIL import Image
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["ai"]["resolution"] = "1k"
old_cover = os.path.join(temp_dir, "old.jpg")
output = os.path.join(temp_dir, "new.jpg")
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
generated = io.BytesIO()
Image.new("RGB", (8, 8), (200, 120, 80)).save(generated, "PNG")
calls = []
downloads = []
events = []
2026-07-07 21:00:43 +08:00
steps = []
2026-07-04 15:13:15 +08:00
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse(
{
"image_url": "https://cdn.example.com/generated.png",
"alias": "image-hd",
"model_used": "provider-image-model",
"points_cost": 8,
"points_balance": 91,
"call_id": "call-image-1",
}
)
def fake_get(url, **kwargs):
downloads.append((url, kwargs))
return _RequestsResponse(content=generated.getvalue())
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
2026-07-04 15:13:15 +08:00
mock.patch(
"app.ai.socket.getaddrinfo",
return_value=[
(
socket.AF_INET,
socket.SOCK_STREAM,
6,
"",
("93.184.216.34", 443),
)
],
):
result = ai.gen_cover(
"生成封面",
old_cover,
output,
resolution="1k",
config=cfg,
cmhub_config_path=key_path,
2026-07-07 21:00:43 +08:00
on_step=steps.append,
2026-07-04 15:13:15 +08:00
on_event=events.append,
)
self.assertEqual(os.path.abspath(output), result)
payload = calls[0][2]["json"]
self.assertEqual("image-hd", payload["model"])
self.assertEqual("1K", payload["resolution"])
self.assertEqual("1:1", payload["aspect_ratio"])
self.assertTrue(payload["image_base64"].startswith("data:image/jpeg;base64,"))
2026-07-08 11:22:47 +08:00
self.assertEqual((3, ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS), calls[0][2]["timeout"])
2026-07-04 15:13:15 +08:00
self.assertEqual("https://cdn.example.com/generated.png", downloads[0][0])
2026-07-08 11:22:47 +08:00
self.assertEqual((3, ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS), downloads[0][1]["timeout"])
2026-07-04 15:13:15 +08:00
self.assertEqual(91, events[0]["metadata"]["points_balance"])
2026-07-07 21:00:43 +08:00
timed_steps = [
event for event in steps
if isinstance(event, dict) and event.get("result") == "success"
]
timed_step_names = [event.get("step") for event in timed_steps]
self.assertIn("cover_request", timed_step_names)
self.assertIn("cover_download", timed_step_names)
self.assertIn("cover_save", timed_step_names)
self.assertTrue(
any(
event.get("step") == "cover_download"
and "下载完成" in event.get("detail", "")
and "耗时" in event.get("detail", "")
for event in timed_steps
)
)
2026-07-04 15:13:15 +08:00
with Image.open(output) as saved:
self.assertEqual((1024, 1024), saved.size)
self.assert_removed(temp_dir)
def test_cmhub_gen_cover_emits_debug_image_url_when_enabled(self):
try:
from PIL import Image
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
old_cover = os.path.join(temp_dir, "old.jpg")
output = os.path.join(temp_dir, "new.jpg")
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
generated_png = self._png_bytes()
steps = []
def fake_request(method, url, **kwargs):
return _RequestsResponse(
{"image_url": "https://cdn.example.com/generated.png?token=secret"}
)
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch.dict(
os.environ,
{"CMSHOPEE_DEBUG_CMHUB_IMAGE_URL": "1"},
), mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch(
"app.ai.socket.getaddrinfo",
return_value=[
(
socket.AF_INET,
socket.SOCK_STREAM,
6,
"",
("93.184.216.34", 443),
)
],
):
ai.gen_cover(
"生成封面",
old_cover,
output,
config=cfg,
cmhub_config_path=key_path,
on_step=steps.append,
)
debug_events = [
event for event in steps
if isinstance(event, dict) and event.get("step") == "cover_image_url"
]
self.assertEqual(1, len(debug_events))
self.assertEqual("debug", debug_events[0]["result"])
self.assertTrue(debug_events[0]["debug_only"])
self.assertIn("https://cdn.example.com/generated.png", debug_events[0]["detail"])
self.assert_removed(temp_dir)
def test_cmhub_gen_cover_retries_image_download_without_new_generation(self):
try:
from PIL import Image
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
old_cover = os.path.join(temp_dir, "old.jpg")
output = os.path.join(temp_dir, "new.jpg")
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
generated_png = self._png_bytes()
request_calls = []
download_calls = []
steps = []
def fake_request(method, url, **kwargs):
request_calls.append((method, url, kwargs))
return _RequestsResponse(
{"image_url": "https://cdn.example.com/generated.png"}
)
def fake_get(url, **kwargs):
download_calls.append((url, kwargs))
if len(download_calls) == 1:
raise ai.requests.exceptions.ConnectionError("temporary")
return _RequestsResponse(content=generated_png)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch(
"app.ai.socket.getaddrinfo",
return_value=[
(
socket.AF_INET,
socket.SOCK_STREAM,
6,
"",
("93.184.216.34", 443),
)
],
), mock.patch("app.ai.time.sleep"):
result = ai.gen_cover(
"生成封面",
old_cover,
output,
config=cfg,
cmhub_config_path=key_path,
on_step=steps.append,
)
self.assertEqual(os.path.abspath(output), result)
self.assertEqual(1, len(request_calls))
self.assertEqual(2, len(download_calls))
retry_events = [
event for event in steps
if isinstance(event, dict)
and event.get("step") == "cover_download"
and event.get("result") == "retry"
]
self.assertEqual(1, len(retry_events))
self.assertEqual(1, retry_events[0]["attempt"])
self.assertEqual(ai.CMHUB_IMAGE_DOWNLOAD_ATTEMPTS, retry_events[0]["attempts"])
self.assert_removed(temp_dir)
def test_cmhub_download_slow_warning_event(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
generated_png = self._png_bytes()
steps = []
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch(
"app.ai.socket.getaddrinfo",
return_value=[
(
socket.AF_INET,
socket.SOCK_STREAM,
6,
"",
("93.184.216.34", 443),
)
],
):
image_bytes, elapsed = ai._download_cmhub_image_with_retry(
"https://cdn.example.com/generated.png",
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
on_step=steps.append,
slow_threshold=0,
)
self.assertEqual(generated_png, image_bytes)
self.assertGreaterEqual(elapsed, 0)
self.assertTrue(
any(
isinstance(event, dict)
and event.get("step") == "cover_download"
and event.get("result") == "warning"
and "图片下载较慢" in event.get("detail", "")
for event in steps
)
)
2026-07-04 15:13:15 +08:00
def test_cmhub_image_url_rejects_private_and_private_dns(self):
with self.assertRaises(ai.AIError):
ai._download_cmhub_image("http://127.0.0.1/a.png", 1, 1)
with mock.patch(
"app.ai.socket.getaddrinfo",
return_value=[
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.2", 443))
],
):
with self.assertRaises(ai.AIError):
ai._download_cmhub_image("https://cdn.example.com/a.png", 1, 1)
def test_cmhub_image_download_uses_curl_without_url_in_argv(self):
generated_png = self._png_bytes()
url = "https://cdn.example.com/generated.png?token=secret-token"
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
calls = []
class FakeProcess:
returncode = 0
def poll(self):
return self.returncode
def communicate(self):
return b"", b""
def fake_popen(args, **kwargs):
calls.append((args, kwargs))
self.assertIn("-K", args)
config_path = args[args.index("-K") + 1]
with open(config_path, "r", encoding="utf-8") as fh:
self.assertIn(url, fh.read())
self.assertNotIn(url, args)
self.assertIn("--connect-timeout", args)
self.assertEqual("3", args[args.index("--connect-timeout") + 1])
self.assertIn("--max-time", args)
2026-07-08 11:22:47 +08:00
self.assertEqual(
str(ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS),
args[args.index("--max-time") + 1],
)
self.assertIn("--max-filesize", args)
self.assertEqual(str(ai.CMHUB_IMAGE_MAX_BYTES), args[args.index("--max-filesize") + 1])
self.assertIn("--noproxy", args)
self.assertEqual("*", args[args.index("--noproxy") + 1])
self.assertFalse(kwargs["shell"])
2026-07-08 17:49:54 +08:00
self.assertEqual(0x08000000, kwargs["creationflags"])
output_path = args[args.index("--output") + 1]
with open(output_path, "wb") as fh:
fh.write(generated_png)
return FakeProcess()
2026-07-08 17:49:54 +08:00
with mock.patch("app.ai.os.name", "nt"), \
mock.patch("app.ai.subprocess.CREATE_NO_WINDOW", 0x08000000, create=True), \
mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch("app.ai.subprocess.Popen", side_effect=fake_popen), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
url,
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
use_system_proxy=False,
download_with_curl="true",
)
self.assertEqual(generated_png, image_bytes)
self.assertEqual(1, len(calls))
2026-07-08 17:49:54 +08:00
def test_cmhub_curl_hidden_window_kwargs_are_windows_only(self):
with mock.patch("app.ai.os.name", "nt"), \
mock.patch("app.ai.subprocess.CREATE_NO_WINDOW", 0x08000000, create=True):
self.assertEqual(
{"creationflags": 0x08000000},
ai._subprocess_hidden_window_kwargs(),
)
with mock.patch("app.ai.os.name", "posix"):
self.assertEqual({}, ai._subprocess_hidden_window_kwargs())
def test_cmhub_image_download_skips_curl_for_private_url(self):
with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch("app.ai.subprocess.Popen") as popen:
with self.assertRaises(ai.AIError):
ai._download_cmhub_image(
"http://127.0.0.1/a.png",
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
download_with_curl="true",
)
popen.assert_not_called()
def test_cmhub_image_download_falls_back_to_requests_when_curl_fails(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
class FailedProcess:
returncode = 28
def poll(self):
return self.returncode
def communicate(self):
return b"", b"timeout"
with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch(
"app.ai.subprocess.Popen",
return_value=FailedProcess(),
) as popen, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get) as get, \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
download_with_curl="true",
)
self.assertEqual(generated_png, image_bytes)
self.assertEqual(1, popen.call_count)
self.assertEqual(1, get.call_count)
def test_cmhub_image_download_auto_without_curl_uses_requests(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch("app.ai.os.name", "nt"), \
mock.patch("app.ai._find_system_curl", return_value=""), \
mock.patch("app.ai.subprocess.Popen") as popen, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
download_with_curl="auto",
)
self.assertEqual(generated_png, image_bytes)
popen.assert_not_called()
def test_cmhub_image_download_auto_on_non_windows_uses_requests(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch("app.ai.os.name", "posix"), \
mock.patch("app.ai._find_system_curl", return_value="/usr/bin/curl"), \
mock.patch("app.ai.subprocess.Popen") as popen, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
2026-07-08 11:22:47 +08:00
read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
download_with_curl="auto",
)
self.assertEqual(generated_png, image_bytes)
popen.assert_not_called()
def test_cmhub_requests_download_cancel_closes_response(self):
stopped = {"value": False}
class StreamingResponse(_RequestsResponse):
def iter_content(self, chunk_size=65536):
yield b"first"
stopped["value"] = True
yield b"second"
response = StreamingResponse()
with mock.patch.object(
ai._cmhub_session(),
"get",
return_value=response,
):
with self.assertRaises(CancelledError):
ai._download_cmhub_image_with_requests(
"https://cdn.example.com/generated.png",
connect_timeout=3,
read_timeout=30,
max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
should_stop=lambda: stopped["value"],
)
self.assertTrue(response.closed)
def test_cmhub_curl_download_cancel_terminates_process_and_cleans_temp_files(self):
stopped = {"value": False}
captured_paths = []
class RunningProcess:
def __init__(self):
self.returncode = None
self.terminated = False
self.killed = False
def poll(self):
return self.returncode
def terminate(self):
self.terminated = True
self.returncode = -15
def kill(self):
self.killed = True
self.returncode = -9
def wait(self, timeout=None):
return self.returncode
def communicate(self):
return b"", b""
process = RunningProcess()
def fake_popen(args, **kwargs):
captured_paths.extend(
[
args[args.index("-K") + 1],
args[args.index("--output") + 1],
]
)
stopped["value"] = True
return process
with mock.patch(
"app.ai._find_system_curl",
return_value=r"C:\Windows\System32\curl.exe",
), mock.patch(
"app.ai.subprocess.Popen",
side_effect=fake_popen,
):
with self.assertRaises(CancelledError):
ai._download_cmhub_image_with_curl(
"https://cdn.example.com/generated.png",
connect_timeout=3,
read_timeout=30,
max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
should_stop=lambda: stopped["value"],
)
self.assertTrue(process.terminated)
self.assertFalse(process.killed)
self.assertTrue(captured_paths)
self.assertTrue(all(not os.path.exists(path) for path in captured_paths))
2026-07-04 15:13:15 +08:00
def test_cmhub_upstream_error_retries_and_keeps_metadata(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
if len(calls) == 1:
return _RequestsResponse(
{"error": {"code": "upstream_error", "message": "bad gateway"}},
status_code=502,
)
return _RequestsResponse({"titles": ["新标题"], "points_balance": 10})
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
2026-07-04 15:13:15 +08:00
mock.patch("app.ai.time.sleep"):
title = ai.gen_title(
"prompt",
"old",
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("新标题", title)
self.assertEqual(2, len(calls))
self.assert_removed(temp_dir)
def test_cmhub_image_read_timeout_does_not_retry(self):
try:
from PIL import Image
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
old_cover = os.path.join(temp_dir, "old.jpg")
output = os.path.join(temp_dir, "new.jpg")
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
raise ai.requests.exceptions.ReadTimeout("slow")
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
2026-07-04 15:13:15 +08:00
with self.assertRaises(ai.CMHubError) as raised:
ai.gen_cover(
"prompt",
old_cover,
output,
retry=3,
config=cfg,
cmhub_config_path=key_path,
)
self.assertEqual("read_timeout", raised.exception.code)
self.assertEqual(1, len(calls))
2026-07-08 11:22:47 +08:00
self.assertEqual((3, ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS), calls[0][2]["timeout"])
2026-07-04 15:13:15 +08:00
self.assert_removed(temp_dir)
def test_fetch_cmhub_models_returns_aliases(self):
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse(
{
"models": [
{
"alias": "title-standard",
"operation_type": "title",
"requires_image": False,
"pricing_status": "priced",
"prices": [{"resolution": "512", "points_cost": 1}],
}
]
}
)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
2026-07-04 15:13:15 +08:00
models = ai.fetch_cmhub_models("https://cmhub.example.com", "sk-cmhub-secret")
self.assertEqual("GET", calls[0][0])
self.assertEqual("https://cmhub.example.com/api/v1/models", calls[0][1])
self.assertEqual("title-standard", models[0]["alias"])
def test_fetch_cmhub_models_404_returns_clear_base_url_error(self):
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse({"detail": "notfound"}, status_code=404)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
with self.assertRaises(ai.CMHubError) as raised:
ai.fetch_cmhub_models(
"https://cmhub.example.com/api/v1/",
"sk-cmhub-secret",
)
self.assertEqual("GET", calls[0][0])
self.assertEqual("https://cmhub.example.com/api/v1/models", calls[0][1])
self.assertEqual("not_found", raised.exception.code)
self.assertEqual(404, raised.exception.status)
message = str(raised.exception)
self.assertIn("cmhub 接口不存在", message)
self.assertIn("/api/v1/models", message)
self.assertNotIn("notfound", message)
2026-07-04 16:29:15 +08:00
def test_fetch_cmhub_balance_returns_points(self):
calls = []
def fake_request(method, url, **kwargs):
calls.append((method, url, kwargs))
return _RequestsResponse({"user": {"id": "u1"}, "points_balance": 42})
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
2026-07-04 16:29:15 +08:00
balance = ai.fetch_cmhub_balance("https://cmhub.example.com", "sk-cmhub-secret")
self.assertEqual("GET", calls[0][0])
self.assertEqual("https://cmhub.example.com/api/v1/balance", calls[0][1])
self.assertEqual(42, balance["points_balance"])
2026-07-07 21:00:43 +08:00
def test_cmhub_image_concurrency_plan_caps_at_five(self):
plan = ai.cmhub_image_concurrency_plan({"image_concurrency": 10})
self.assertEqual(10, plan["configured_image_concurrency"])
self.assertEqual(5, plan["request_concurrency"])
self.assertEqual(5, plan["download_concurrency"])
self.assertEqual(5, plan["limit"])
def test_generate_batch_cmhub_image_downloads_do_not_block_later_requests(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
2026-07-09 17:44:03 +08:00
cfg["ai"]["generate_mode"] = "cover"
cfg["ai"]["generate_mode"] = "cover"
2026-07-07 21:00:43 +08:00
cfg["ai"]["generate_cover"] = True
cfg["ai"]["image_concurrency"] = 10
titles = ["旧标题%s" % index for index in range(7)]
batch_id, tasks = self._collected_tasks(temp_dir, cfg, titles)
self._write_old_cover_files(tasks)
for index, task in enumerate(tasks):
db.set_generated(task.id, "已有标题%s" % index, None, path=cfg["db_path"])
tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
lock = threading.Lock()
all_requests_seen = threading.Event()
counters = {"request_count": 0, "active": 0, "max_active": 0}
downloads = []
def fake_request(method, url, **kwargs):
2026-07-09 00:07:09 +08:00
if str(method).upper() == "GET":
task_id = url.rsplit("/", 1)[-1]
index = int(task_id.rsplit("-", 1)[-1])
return _RequestsResponse(
{
"task_id": task_id,
"status": "succeeded",
"result": {
"image_url": "https://cdn.example.com/generated-%s.png"
% index
},
}
)
2026-07-07 21:00:43 +08:00
with lock:
counters["request_count"] += 1
request_index = counters["request_count"]
counters["active"] += 1
counters["max_active"] = max(
counters["max_active"],
counters["active"],
)
if counters["request_count"] >= len(tasks):
all_requests_seen.set()
try:
return _RequestsResponse(
{
2026-07-09 00:07:09 +08:00
"task_id": "cmhub-task-%s" % request_index,
"status": "queued",
2026-07-07 21:00:43 +08:00
}
)
finally:
with lock:
counters["active"] -= 1
def fake_get(url, **kwargs):
self.assertTrue(
all_requests_seen.wait(2),
"慢下载不能阻塞后续 cmhub 生图请求提交",
)
with lock:
downloads.append((url, kwargs))
return _RequestsResponse(content=generated_png)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
2026-07-07 21:00:43 +08:00
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
},
)
self.assertTrue(summary["ok"])
self.assertEqual(len(tasks), summary["cover_done"])
self.assertEqual(len(tasks), summary["generated_done"])
self.assertEqual(len(tasks), counters["request_count"])
self.assertLessEqual(counters["max_active"], 5)
self.assertEqual(len(tasks), len(downloads))
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
self.assertTrue(all(os.path.exists(task.new_cover_path) for task in updated))
self.assert_removed(temp_dir)
def test_generate_batch_cmhub_download_failure_does_not_request_image_again(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
cfg["ai"]["image_concurrency"] = 10
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
self._write_old_cover_files(tasks)
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
requests_seen = []
downloads_seen = []
def fake_request(method, url, **kwargs):
2026-07-09 00:07:09 +08:00
if str(method).upper() == "GET":
return _RequestsResponse(
{
"task_id": "cmhub-task-1",
"status": "succeeded",
"result": {"image_url": "https://cdn.example.com/generated.png"},
}
)
2026-07-07 21:00:43 +08:00
requests_seen.append((method, url, kwargs))
2026-07-09 00:07:09 +08:00
return _RequestsResponse({"task_id": "cmhub-task-1", "status": "queued"})
2026-07-07 21:00:43 +08:00
def fake_get(url, **kwargs):
downloads_seen.append((url, kwargs))
raise ai.requests.exceptions.ConnectionError("download failed")
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
2026-07-07 21:00:43 +08:00
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
},
)
self.assertFalse(summary["ok"])
self.assertEqual(0, summary["cover_done"])
self.assertEqual(1, summary["failed"])
self.assertEqual(1, len(requests_seen))
self.assertEqual(ai.CMHUB_IMAGE_DOWNLOAD_ATTEMPTS, len(downloads_seen))
2026-07-07 21:00:43 +08:00
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
self.assertEqual("failed", updated.status)
self.assertIn("下载 cmhub 图片失败", updated.last_error)
self.assert_removed(temp_dir)
2026-07-09 00:07:09 +08:00
def test_generate_batch_cmhub_async_submit_persists_task_id_before_poll(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
self._write_old_cover_files(tasks)
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
poll_db_values = []
headers_seen = []
2026-07-09 14:30:53 +08:00
submit_timeouts = []
2026-07-09 08:57:46 +08:00
downloads = []
2026-07-09 00:07:09 +08:00
def fake_request(method, url, **kwargs):
if str(method).upper() == "POST":
headers_seen.append(dict(kwargs["headers"]))
2026-07-09 14:30:53 +08:00
submit_timeouts.append(kwargs.get("timeout"))
2026-07-09 00:07:09 +08:00
return _RequestsResponse(
{
"task_id": "cmhub-task-1",
"status": "queued",
"points_cost": 2,
"points_balance": 80,
"call_id": "call-image-1",
},
status_code=202,
)
poll_db_values.append(db.get_task(task.id, path=cfg["db_path"]).image_task_id)
return _RequestsResponse(
{
"task_id": "cmhub-task-1",
"status": "succeeded",
2026-07-09 08:57:46 +08:00
"result": {"image_url": "/generated/images/generated.png"},
2026-07-09 00:07:09 +08:00
}
)
def fake_get(url, **kwargs):
2026-07-09 08:57:46 +08:00
downloads.append((url, kwargs))
2026-07-09 00:07:09 +08:00
return _RequestsResponse(content=generated_png)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
summary = ai.generate_batch(
[task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
},
)
self.assertTrue(summary["ok"])
self.assertEqual(["cmhub-task-1"], poll_db_values)
self.assertTrue(headers_seen[0]["Idempotency-Key"].startswith("cmshopee-task-"))
self.assertTrue(headers_seen[0]["X-Client-Version"])
2026-07-09 14:30:53 +08:00
self.assertEqual([(3, 36)], submit_timeouts)
2026-07-09 08:57:46 +08:00
self.assertEqual(
"https://cmhub.example.com/generated/images/generated.png",
downloads[0][0],
)
2026-07-09 00:07:09 +08:00
updated = db.get_task(task.id, path=cfg["db_path"])
self.assertEqual("cmhub-task-1", updated.image_task_id)
self.assertTrue(updated.image_task_key)
self.assertTrue(os.path.exists(updated.new_cover_path))
self.assert_removed(temp_dir)
2026-07-09 08:57:46 +08:00
def test_cmhub_async_image_url_extractor_accepts_image_urls_list(self):
data = {
"task_id": "cmhub-task-1",
"status": "succeeded",
"result": {
"image_urls": [
"https://cdn.example.com/generated-a.png",
"https://cdn.example.com/generated-b.png",
]
},
}
self.assertEqual(
"https://cdn.example.com/generated-a.png",
ai._extract_cmhub_image_url(data, "https://cmhub.example.com"),
)
2026-07-09 00:07:09 +08:00
def test_generate_batch_cmhub_resumes_existing_image_task_without_submit(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
self._write_old_cover_files(tasks)
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
key = db.ensure_image_task_key(tasks[0].id, path=cfg["db_path"])
db.set_image_task_submitted(tasks[0].id, "cmhub-task-resume", key, path=cfg["db_path"])
task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
posts = []
def fake_request(method, url, **kwargs):
if str(method).upper() == "POST":
posts.append((method, url, kwargs))
return _RequestsResponse({"task_id": "unexpected", "status": "queued"})
return _RequestsResponse(
{
"task_id": "cmhub-task-resume",
"status": "succeeded",
"result": {"image_url": "https://cdn.example.com/resume.png"},
}
)
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
summary = ai.generate_batch(
[task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
},
)
self.assertTrue(summary["ok"])
self.assertEqual([], posts)
updated = db.get_task(task.id, path=cfg["db_path"])
self.assertEqual("cmhub-task-resume", updated.image_task_id)
self.assertEqual(key, updated.image_task_key)
self.assertTrue(os.path.exists(updated.new_cover_path))
self.assert_removed(temp_dir)
def test_generate_batch_cmhub_failed_task_clears_image_task_state(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
self._write_old_cover_files(tasks)
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
key = db.ensure_image_task_key(tasks[0].id, path=cfg["db_path"])
db.set_image_task_submitted(tasks[0].id, "cmhub-task-failed", key, path=cfg["db_path"])
task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
def fake_request(method, url, **kwargs):
return _RequestsResponse(
{
"task_id": "cmhub-task-failed",
"status": "failed",
"error": {"code": "upstream_timeout", "message": "上游超时"},
}
)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
summary = ai.generate_batch(
[task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
},
)
self.assertFalse(summary["ok"])
self.assertEqual(1, summary["failed"])
updated = db.get_task(task.id, path=cfg["db_path"])
self.assertIsNone(updated.image_task_id)
self.assertIsNone(updated.image_task_key)
self.assertIn("上游生成超时", updated.last_error)
self.assert_removed(temp_dir)
def test_generate_batch_cmhub_cancelled_poll_keeps_image_task_state(self):
try:
from PIL import Image # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
self._write_old_cover_files(tasks)
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
events = []
def fake_request(method, url, **kwargs):
if str(method).upper() == "POST":
return _RequestsResponse({"task_id": "cmhub-task-cancel", "status": "queued"})
return _RequestsResponse({"task_id": "cmhub-task-cancel", "status": "running"})
def should_stop():
current = db.get_task(task.id, path=cfg["db_path"])
return bool(current and current.image_task_id)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
summary = ai.generate_batch(
[task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
"on_event": events.append,
},
should_stop=should_stop,
)
self.assertFalse(summary["ok"])
self.assertTrue(summary["cancelled"])
updated = db.get_task(task.id, path=cfg["db_path"])
self.assertEqual("cmhub-task-cancel", updated.image_task_id)
self.assertTrue(updated.image_task_key)
self.assertIsNone(updated.new_cover_path)
self.assertTrue(
any("服务端任务可能仍在完成" in event.get("detail", "") for event in events)
)
self.assert_removed(temp_dir)
2026-07-07 21:00:43 +08:00
2026-07-04 15:13:15 +08:00
def test_generate_batch_forwards_cmhub_metadata_event(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题A"])
events = []
def fake_request(method, url, **kwargs):
return _RequestsResponse(
{
"titles": ["新标题A"],
"alias": "title-standard",
"points_cost": 1,
"points_balance": 88,
"call_id": "call-batch-1",
}
)
with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
2026-07-04 15:13:15 +08:00
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"cmhub_config_path": key_path,
"on_event": events.append,
},
)
self.assertTrue(summary["ok"])
self.assertTrue(
any(
event.get("metadata", {}).get("points_balance") == 88
and event.get("metadata", {}).get("call_id") == "call-batch-1"
for event in events
)
)
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
self.assertEqual("generated", updated[0].stage)
self.assert_removed(temp_dir)
2026-06-27 16:26:08 +08:00
def test_generate_batch_persists_titles_and_covers_per_task(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["title_concurrency"] = 2
cfg["ai"]["image_concurrency"] = 2
2026-07-02 15:17:11 +08:00
cfg["ai"]["generate_cover"] = True
2026-06-27 16:26:08 +08:00
batch_id, tasks = self._collected_tasks(temp_dir, cfg)
cover_prompts = []
progress = []
def fake_title(title_prompt, old_title, **kwargs):
self.assertEqual("标题提示", title_prompt)
return "新" + old_title
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
cover_prompts.append(cover_prompt)
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title", side_effect=fake_title), \
mock.patch("app.ai.gen_cover", side_effect=fake_cover):
summary = ai.generate_batch(
tasks,
{
"title": "标题提示",
"cover": "封面 {新标题} {店铺} {商品id}",
},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"account_by_alias": {
"alias-a": SimpleNamespace(account_name="主店", slug="main")
},
},
on_progress=progress.append,
)
self.assertTrue(summary["ok"])
self.assertEqual(2, summary["title_done"])
self.assertEqual(2, summary["cover_done"])
self.assertEqual(0, summary["failed"])
self.assertTrue(progress)
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
self.assertTrue(all(task.stage == "generated" for task in updated))
self.assertEqual({"新旧标题A", "新旧标题B"}, {task.new_title for task in updated})
2026-07-01 08:35:46 +08:00
expected_cover_paths = {
os.path.abspath(
os.path.join(
cfg["image_dir"],
str(task.batch_id),
"main",
f"{task.id}_{task.item_id}_new.jpg",
)
)
for task in updated
}
self.assertEqual(expected_cover_paths, {task.new_cover_path for task in updated})
2026-06-27 16:26:08 +08:00
self.assertTrue(all(os.path.exists(task.new_cover_path) for task in updated))
self.assertIn("主店", "\n".join(cover_prompts))
self.assertIn("新旧标题A", "\n".join(cover_prompts))
self.assert_removed(temp_dir)
2026-07-02 15:17:11 +08:00
def test_generate_batch_can_skip_cover_generation(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
batch_id, tasks = self._collected_tasks(temp_dir, cfg)
events = []
progress = []
def fake_title(title_prompt, old_title, **kwargs):
self.assertEqual("标题提示", title_prompt)
return "新" + old_title
with mock.patch("app.ai.gen_title", side_effect=fake_title), \
mock.patch("app.ai.gen_cover") as gen_cover:
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"on_event": events.append,
},
on_progress=progress.append,
)
self.assertTrue(summary["ok"])
self.assertFalse(summary["generate_cover"])
self.assertEqual(2, summary["title_done"])
self.assertEqual(0, summary["cover_done"])
self.assertEqual(0, summary["cover_total"])
self.assertEqual(2, summary["generated_done"])
self.assertEqual(0, summary["failed"])
gen_cover.assert_not_called()
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
self.assertTrue(all(task.stage == "generated" for task in updated))
self.assertTrue(all(task.status == "success" for task in updated))
self.assertEqual({"新旧标题A", "新旧标题B"}, {task.new_title for task in updated})
self.assertTrue(all(task.new_cover_path is None for task in updated))
self.assertEqual(0, progress[-1]["cover_total"])
self.assertEqual(2, progress[-1]["generated_done"])
self.assertTrue(
any(
event.get("phase") == "title"
and event.get("step") == "db_write"
and event.get("result") == "success"
and event.get("detail") == "仅生成标题"
for event in events
)
)
self.assert_removed(temp_dir)
2026-06-27 16:26:08 +08:00
def test_generate_batch_marks_failed_task_without_blocking_others(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
2026-07-02 15:17:11 +08:00
cfg["ai"]["generate_cover"] = True
2026-06-27 16:26:08 +08:00
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["好标题", "坏标题"])
def fake_title(title_prompt, old_title, **kwargs):
if old_title == "坏标题":
raise ai.AIError("标题生成失败")
return "新" + old_title
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title", side_effect=fake_title), \
mock.patch("app.ai.gen_cover", side_effect=fake_cover):
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertFalse(summary["ok"])
self.assertEqual(1, summary["title_done"])
self.assertEqual(1, summary["cover_done"])
self.assertEqual(1, summary["failed"])
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
by_title = {task.old_title: task for task in updated}
self.assertEqual("generated", by_title["好标题"].stage)
self.assertEqual("success", by_title["好标题"].status)
self.assertEqual("collected", by_title["坏标题"].stage)
self.assertEqual("failed", by_title["坏标题"].status)
self.assertIn("标题生成失败", by_title["坏标题"].last_error)
self.assert_removed(temp_dir)
2026-07-06 17:57:39 +08:00
def test_generate_batch_fills_missing_cover_without_regenerating_title(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
2026-07-09 17:44:03 +08:00
cfg["ai"]["generate_mode"] = "cover"
2026-07-06 17:57:39 +08:00
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
db.set_generated(tasks[0].id, "手动标题", None, path=cfg["db_path"])
cover_only_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
progress = []
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
summary = ai.generate_batch(
[cover_only_task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
on_progress=progress.append,
)
self.assertTrue(summary["ok"])
2026-07-09 17:44:03 +08:00
self.assertEqual("cover", summary["generate_mode"])
self.assertTrue(summary["generate_cover"])
2026-07-06 17:57:39 +08:00
self.assertEqual(1, summary["total"])
self.assertEqual(0, summary["title_total"])
self.assertEqual(0, summary["title_done"])
self.assertEqual(1, summary["cover_total"])
self.assertEqual(1, summary["cover_done"])
self.assertEqual(1, summary["generated_done"])
self.assertEqual(0, progress[-1]["title_total"])
self.assertEqual(1, progress[-1]["cover_total"])
gen_title.assert_not_called()
gen_cover.assert_called_once()
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
self.assertEqual("generated", updated.stage)
self.assertEqual("success", updated.status)
self.assertEqual("手动标题", updated.new_title)
self.assertTrue(os.path.exists(updated.new_cover_path))
self.assert_removed(temp_dir)
def test_generate_batch_cover_only_uses_old_title_and_keeps_new_title_empty(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_mode"] = "cover"
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
cover_prompts = []
events = []
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
cover_prompts.append(cover_prompt)
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面 {新标题} / {旧标题}"},
ai_cfg={
"config": cfg,
"db_path": cfg["db_path"],
"on_event": events.append,
},
)
self.assertTrue(summary["ok"])
self.assertEqual(0, summary["title_total"])
self.assertEqual(1, summary["cover_total"])
gen_title.assert_not_called()
gen_cover.assert_called_once()
self.assertEqual(["封面 旧标题 / 旧标题"], cover_prompts)
updated = db.get_task(tasks[0].id, path=cfg["db_path"])
self.assertIsNone(updated.new_title)
self.assertTrue(os.path.exists(updated.new_cover_path))
self.assertTrue(
any("使用旧标题作为封面参考" in str(event.get("detail") or "") for event in events)
)
cfg["ai"]["generate_mode"] = "title"
cfg["ai"]["generate_cover"] = False
with mock.patch("app.ai.gen_title", return_value="后补新标题") as gen_title, \
mock.patch("app.ai.gen_cover") as gen_cover:
title_summary = ai.generate_batch(
[updated],
{"title": "标题提示", "cover": "封面"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(title_summary["ok"])
gen_title.assert_called_once()
gen_cover.assert_not_called()
completed = db.get_task(tasks[0].id, path=cfg["db_path"])
self.assertEqual("后补新标题", completed.new_title)
self.assertEqual(updated.new_cover_path, completed.new_cover_path)
self.assert_removed(temp_dir)
def test_generate_batch_fills_reset_cover_after_committed_history(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
db.set_generated(tasks[0].id, "手动标题", "old-new.jpg", path=cfg["db_path"])
db.set_applied(tasks[0].id, True, path=cfg["db_path"])
db.reset_generated(
tasks[0].id,
reset_title=False,
reset_cover=True,
path=cfg["db_path"],
)
cover_only_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
summary = ai.generate_batch(
[cover_only_task],
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(summary["ok"])
self.assertEqual(1, summary["total"])
self.assertEqual(0, summary["title_total"])
self.assertEqual(1, summary["cover_total"])
gen_title.assert_not_called()
gen_cover.assert_called_once()
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
self.assertEqual("手动标题", updated.new_title)
self.assertTrue(os.path.exists(updated.new_cover_path))
self.assertEqual(1, updated.committed)
self.assertEqual(1, updated.apply_attempts)
self.assert_removed(temp_dir)
2026-07-06 17:57:39 +08:00
def test_generate_batch_handles_mixed_title_and_cover_gaps(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题A", "旧标题B"])
db.set_generated(tasks[1].id, "已有标题B", None, path=cfg["db_path"])
mixed_tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
cover_prompts = []
def fake_title(title_prompt, old_title, **kwargs):
return "新" + old_title
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
cover_prompts.append(cover_prompt)
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "wb") as fh:
fh.write(b"jpeg")
return out_path
with mock.patch("app.ai.gen_title", side_effect=fake_title) as gen_title, \
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
summary = ai.generate_batch(
mixed_tasks,
{"title": "标题提示", "cover": "封面 {新标题}"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(summary["ok"])
self.assertEqual(2, summary["total"])
self.assertEqual(1, summary["title_total"])
self.assertEqual(1, summary["title_done"])
self.assertEqual(2, summary["cover_total"])
self.assertEqual(2, summary["cover_done"])
self.assertEqual(2, summary["generated_done"])
self.assertEqual(1, gen_title.call_count)
self.assertEqual(2, gen_cover.call_count)
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
by_old_title = {task.old_title: task for task in updated}
self.assertEqual("新旧标题A", by_old_title["旧标题A"].new_title)
self.assertEqual("已有标题B", by_old_title["旧标题B"].new_title)
self.assertTrue(all(os.path.exists(task.new_cover_path) for task in updated))
self.assertIn("新旧标题A", "\n".join(cover_prompts))
self.assertIn("已有标题B", "\n".join(cover_prompts))
self.assert_removed(temp_dir)
def test_generate_batch_ignores_title_only_task_when_cover_disabled(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = False
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
2026-07-06 17:57:39 +08:00
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
title_only_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
2026-07-06 17:57:39 +08:00
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover") as gen_cover:
summary = ai.generate_batch(
2026-07-06 17:57:39 +08:00
[title_only_task],
{"title": "标题提示", "cover": "封面"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(summary["ok"])
2026-07-06 17:57:39 +08:00
self.assertEqual(0, summary["total"])
self.assertEqual(0, summary["title_total"])
self.assertEqual(0, summary["cover_total"])
gen_title.assert_not_called()
gen_cover.assert_not_called()
self.assert_removed(temp_dir)
def test_generate_batch_ignores_complete_generated_task(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
cfg["ai"]["generate_cover"] = True
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
db.set_generated(tasks[0].id, "已有标题", "new.jpg", path=cfg["db_path"])
complete_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover") as gen_cover:
summary = ai.generate_batch(
[complete_task],
{"title": "标题提示", "cover": "封面"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(summary["ok"])
self.assertEqual(0, summary["total"])
self.assertEqual(0, summary["title_total"])
self.assertEqual(0, summary["cover_total"])
gen_title.assert_not_called()
gen_cover.assert_not_called()
self.assert_removed(temp_dir)
def test_generate_batch_does_not_retry_apply_failed_records(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
db.set_generated(tasks[0].id, "新标题", "new.jpg", path=cfg["db_path"])
db.mark_failed(tasks[0].id, "apply", "更新失败", path=cfg["db_path"])
update_failed_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover") as gen_cover:
summary = ai.generate_batch(
[update_failed_task],
{"title": "标题提示", "cover": "封面"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
)
self.assertTrue(summary["ok"])
self.assertEqual(0, summary["total"])
gen_title.assert_not_called()
gen_cover.assert_not_called()
unchanged = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
self.assertEqual("generated", unchanged.stage)
self.assertEqual("failed", unchanged.status)
self.assertEqual(1, unchanged.apply_attempts)
self.assert_removed(temp_dir)
2026-06-27 16:26:08 +08:00
def test_generate_batch_stop_before_scheduling_keeps_tasks_collected(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config()
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
cfg["image_dir"] = os.path.join(temp_dir, "images")
batch_id, tasks = self._collected_tasks(temp_dir, cfg)
with mock.patch("app.ai.gen_title") as gen_title, \
mock.patch("app.ai.gen_cover") as gen_cover:
summary = ai.generate_batch(
tasks,
{"title": "标题提示", "cover": "封面"},
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
should_stop=lambda: True,
)
self.assertFalse(summary["ok"])
self.assertTrue(summary["cancelled"])
self.assertEqual(0, summary["title_done"])
self.assertEqual(0, summary["cover_done"])
self.assertEqual(0, summary["failed"])
gen_title.assert_not_called()
gen_cover.assert_not_called()
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
self.assertTrue(all(task.stage == "collected" for task in updated))
self.assert_removed(temp_dir)
def test_cmhub_session_is_shared_singleton_with_pool(self):
session_a = ai._cmhub_session()
session_b = ai._cmhub_session()
self.assertIs(session_a, session_b)
# 连接池要覆盖生图并发 + 下载并发(各上限 5)
self.assertGreaterEqual(
ai.CMHUB_HTTP_POOL_SIZE, ai.CMHUB_IMAGE_CONCURRENCY_LIMIT * 2
)
adapter = session_a.get_adapter("https://cmhub.example.com")
self.assertEqual(ai.CMHUB_HTTP_POOL_SIZE, adapter._pool_maxsize)
def test_apply_cmhub_proxy_toggles_trust_env(self):
try:
session = ai._apply_cmhub_proxy(False)
self.assertIs(session, ai._cmhub_session())
self.assertFalse(session.trust_env)
ai._apply_cmhub_proxy(True)
self.assertTrue(ai._cmhub_session().trust_env)
finally:
# 复位为默认(绕过系统代理),避免影响其它测试
ai._apply_cmhub_proxy(False)
def test_cmhub_runtime_applies_and_returns_use_system_proxy(self):
with mock.patch(
"app.ai.appconfig.cmhub_config",
return_value={
"base_url": "https://cmhub.example.com",
"title_alias": "title-standard",
"image_alias": "image-hd",
"connect_timeout": 10,
"use_system_proxy": True,
},
), mock.patch(
"app.ai.appconfig.get_cmhub_api_key", return_value="sk_cmhub_test"
):
try:
runtime = ai._cmhub_runtime({}, "image", cmhub_config_path=None)
self.assertTrue(runtime["use_system_proxy"])
self.assertTrue(ai._cmhub_session().trust_env)
finally:
ai._apply_cmhub_proxy(False)
2026-06-27 15:47:21 +08:00
if __name__ == "__main__":
unittest.main()