1196 lines
50 KiB
Python
1196 lines
50 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import socket
|
|
import sys
|
|
import threading
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from _helpers import TempDirMixin
|
|
|
|
from app import ai, appconfig, db
|
|
|
|
|
|
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")
|
|
|
|
|
|
|
|
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)
|
|
|
|
def json(self):
|
|
return self.payload
|
|
|
|
def iter_content(self, chunk_size=65536):
|
|
if self.content:
|
|
yield self.content
|
|
|
|
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()
|
|
cfg["ai"]["backend"] = "direct"
|
|
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
|
|
|
|
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",
|
|
"connect_timeout": 3,
|
|
"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
|
|
|
|
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 _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()
|
|
|
|
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 = []
|
|
|
|
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,
|
|
)
|
|
|
|
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"])
|
|
body = json.loads(calls[-1][0].data.decode("utf-8"))
|
|
self.assertEqual("text-model", body["model"])
|
|
self.assertEqual(0, body["temperature"])
|
|
self.assertNotIn("sk-text-secret", body["messages"][1]["content"])
|
|
|
|
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)
|
|
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)
|
|
|
|
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("app.ai.requests.request", side_effect=fake_request):
|
|
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)
|
|
self.assertEqual((3, 600), kwargs["timeout"])
|
|
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_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))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
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 = []
|
|
steps = []
|
|
|
|
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("app.ai.requests.request", side_effect=fake_request), \
|
|
mock.patch("app.ai.requests.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),
|
|
)
|
|
],
|
|
):
|
|
result = ai.gen_cover(
|
|
"生成封面",
|
|
old_cover,
|
|
output,
|
|
resolution="1k",
|
|
config=cfg,
|
|
cmhub_config_path=key_path,
|
|
on_step=steps.append,
|
|
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,"))
|
|
self.assertEqual((3, 650), calls[0][2]["timeout"])
|
|
self.assertEqual("https://cdn.example.com/generated.png", downloads[0][0])
|
|
self.assertEqual((3, 650), downloads[0][1]["timeout"])
|
|
self.assertEqual(91, events[0]["metadata"]["points_balance"])
|
|
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
|
|
)
|
|
)
|
|
with Image.open(output) as saved:
|
|
self.assertEqual((1024, 1024), saved.size)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
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_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("app.ai.requests.request", side_effect=fake_request), \
|
|
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("app.ai.requests.request", side_effect=fake_request):
|
|
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))
|
|
self.assertEqual((3, 650), calls[0][2]["timeout"])
|
|
|
|
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("app.ai.requests.request", side_effect=fake_request):
|
|
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("app.ai.requests.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)
|
|
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("app.ai.requests.request", side_effect=fake_request):
|
|
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"])
|
|
|
|
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")
|
|
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):
|
|
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(
|
|
{
|
|
"image_url": "https://cdn.example.com/generated-%s.png"
|
|
% request_index
|
|
}
|
|
)
|
|
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("app.ai.requests.request", side_effect=fake_request), \
|
|
mock.patch("app.ai.requests.get", side_effect=fake_get), \
|
|
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):
|
|
requests_seen.append((method, url, kwargs))
|
|
return _RequestsResponse(
|
|
{"image_url": "https://cdn.example.com/generated.png"}
|
|
)
|
|
|
|
def fake_get(url, **kwargs):
|
|
downloads_seen.append((url, kwargs))
|
|
raise ai.requests.exceptions.ConnectionError("download failed")
|
|
|
|
with mock.patch("app.ai.requests.request", side_effect=fake_request), \
|
|
mock.patch("app.ai.requests.get", side_effect=fake_get), \
|
|
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(1, len(downloads_seen))
|
|
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)
|
|
|
|
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("app.ai.requests.request", side_effect=fake_request):
|
|
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)
|
|
|
|
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
|
|
cfg["ai"]["generate_cover"] = True
|
|
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})
|
|
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})
|
|
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)
|
|
|
|
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)
|
|
|
|
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")
|
|
cfg["ai"]["generate_cover"] = True
|
|
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)
|
|
|
|
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")
|
|
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"])
|
|
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_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)
|
|
|
|
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, ["旧标题"])
|
|
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]
|
|
|
|
with mock.patch("app.ai.gen_title") as gen_title, \
|
|
mock.patch("app.ai.gen_cover") as gen_cover:
|
|
summary = ai.generate_batch(
|
|
[title_only_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_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)
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|