326 lines
11 KiB
Python
326 lines
11 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
|
|
|
|
from _helpers import REPO_ROOT # noqa: F401
|
|
|
|
from app import db, image_studio, image_studio_images, workers
|
|
|
|
if workers.QT_IMPORT_ERROR is not None:
|
|
raise unittest.SkipTest("PySide6 未安装")
|
|
|
|
from PySide6.QtCore import QEventLoop, QTimer
|
|
from PySide6.QtWidgets import QApplication
|
|
|
|
from app.workers import BaseWorker, run_worker
|
|
from app.gui.workers import (
|
|
ImageStudioDownloadOriginalWorker,
|
|
ImageStudioPullImagesWorker,
|
|
ProductSuiteAiWriteWorker,
|
|
ProductSuiteGenerateWorker,
|
|
ProductSuiteHistoryExportWorker,
|
|
)
|
|
|
|
|
|
class DemoWorker(BaseWorker):
|
|
def execute(self):
|
|
self.log.emit("开始")
|
|
self.progress.emit({"done": 1, "total": 1})
|
|
self.row_updated.emit(7, {"status": "success"})
|
|
return {"ok": True, "done": 1}
|
|
|
|
|
|
class CancelAwareWorker(BaseWorker):
|
|
def execute(self):
|
|
if self.should_cancel():
|
|
return {"done": 0}
|
|
self.cancel()
|
|
return {"done": 0}
|
|
|
|
|
|
class FailingWorker(BaseWorker):
|
|
def execute(self):
|
|
raise RuntimeError("模拟失败")
|
|
|
|
|
|
class WorkerTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.app = QApplication.instance() or QApplication([])
|
|
|
|
def start_and_wait(self, thread, timeout_ms=2000):
|
|
loop = QEventLoop()
|
|
finished = []
|
|
|
|
def on_finished():
|
|
finished.append(True)
|
|
loop.quit()
|
|
|
|
timer = QTimer()
|
|
timer.setSingleShot(True)
|
|
timer.timeout.connect(loop.quit)
|
|
thread.finished.connect(on_finished)
|
|
timer.start(timeout_ms)
|
|
thread.start()
|
|
loop.exec()
|
|
|
|
self.assertTrue(finished, "worker thread did not finish before timeout")
|
|
|
|
def test_worker_emits_common_signals_and_finishes(self):
|
|
worker = DemoWorker()
|
|
logs = []
|
|
progress = []
|
|
rows = []
|
|
finished = []
|
|
|
|
worker.log.connect(logs.append)
|
|
worker.progress.connect(lambda payload: progress.append(dict(payload)))
|
|
worker.row_updated.connect(lambda task_id, fields: rows.append((task_id, dict(fields))))
|
|
worker.finished.connect(lambda payload: finished.append(dict(payload)))
|
|
|
|
thread = run_worker(worker, start=False)
|
|
self.start_and_wait(thread)
|
|
|
|
self.assertEqual(["开始"], logs)
|
|
self.assertEqual([{"done": 1, "total": 1}], progress)
|
|
self.assertEqual([(7, {"status": "success"})], rows)
|
|
self.assertEqual([{"ok": True, "done": 1}], finished)
|
|
|
|
def test_cancel_flag_emits_cancelled_instead_of_finished(self):
|
|
worker = CancelAwareWorker()
|
|
cancelled = []
|
|
finished = []
|
|
|
|
worker.cancelled.connect(lambda payload: cancelled.append(dict(payload)))
|
|
worker.finished.connect(lambda payload: finished.append(dict(payload)))
|
|
|
|
thread = run_worker(worker, start=False)
|
|
self.start_and_wait(thread)
|
|
|
|
self.assertEqual([{"done": 0, "cancelled": True}], cancelled)
|
|
self.assertEqual([], finished)
|
|
|
|
def test_uncaught_exception_emits_failed_and_finished_summary(self):
|
|
worker = FailingWorker()
|
|
failed = []
|
|
finished = []
|
|
|
|
worker.failed.connect(lambda task_id, error: failed.append((task_id, error)))
|
|
worker.finished.connect(lambda payload: finished.append(dict(payload)))
|
|
|
|
thread = run_worker(worker, start=False)
|
|
self.start_and_wait(thread)
|
|
|
|
self.assertEqual([(-1, "模拟失败")], failed)
|
|
self.assertEqual([{"ok": False, "error": "模拟失败"}], finished)
|
|
|
|
def test_run_worker_rejects_plain_object(self):
|
|
with self.assertRaises(TypeError):
|
|
run_worker(object(), start=False)
|
|
|
|
def test_image_studio_original_download_retries_twice_without_raw_error(self):
|
|
worker = ImageStudioDownloadOriginalWorker(
|
|
12,
|
|
max_retries=2,
|
|
retry_delays=(0, 0),
|
|
)
|
|
progress = []
|
|
logs = []
|
|
worker.progress.connect(lambda payload: progress.append(dict(payload)))
|
|
worker.log.connect(logs.append)
|
|
successful_asset = SimpleNamespace(id=12)
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio_images.download_original_asset",
|
|
side_effect=[RuntimeError("first"), RuntimeError("second"), successful_asset],
|
|
) as download:
|
|
summary = worker.execute()
|
|
|
|
self.assertEqual(3, download.call_count)
|
|
self.assertIs(successful_asset, summary["asset"])
|
|
retries = [payload for payload in progress if payload.get("state") == "retry"]
|
|
self.assertEqual(
|
|
[
|
|
{"asset_id": 12, "state": "retry", "retry": 1, "max_retries": 2, "delay_seconds": 0.0},
|
|
{"asset_id": 12, "state": "retry", "retry": 2, "max_retries": 2, "delay_seconds": 0.0},
|
|
],
|
|
retries,
|
|
)
|
|
self.assertTrue(any("重试 1/2" in message for message in logs))
|
|
self.assertTrue(any("重试 2/2" in message for message in logs))
|
|
self.assertFalse(any("first" in message or "second" in message for message in logs))
|
|
|
|
def test_image_studio_original_download_returns_chinese_final_failure(self):
|
|
worker = ImageStudioDownloadOriginalWorker(
|
|
12,
|
|
max_retries=2,
|
|
retry_delays=(0, 0),
|
|
)
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio_images.download_original_asset",
|
|
side_effect=RuntimeError("https://example.invalid/private"),
|
|
) as download:
|
|
summary = worker.execute()
|
|
|
|
self.assertEqual(3, download.call_count)
|
|
self.assertFalse(summary["ok"])
|
|
self.assertEqual("蝦皮原主图下载失败，请稍后再次点击图片重试。", summary["error"])
|
|
self.assertNotIn("https://", summary["error"])
|
|
|
|
def test_image_studio_pull_worker_returns_token_when_cancelled(self):
|
|
worker = ImageStudioPullImagesWorker(
|
|
"alias-a",
|
|
"51100639510",
|
|
pull_run_token="pull-token",
|
|
)
|
|
worker.cancel()
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio.pull_remote_main_image_urls",
|
|
) as pull:
|
|
summary = worker.execute()
|
|
|
|
pull.assert_not_called()
|
|
self.assertTrue(summary["cancelled"])
|
|
self.assertEqual("pull-token", summary["pull_run_token"])
|
|
|
|
def test_image_studio_original_download_maps_safe_cancel(self):
|
|
worker = ImageStudioDownloadOriginalWorker(12)
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio_images.download_original_asset",
|
|
side_effect=image_studio_images.ImageStudioImageCancelled("停止"),
|
|
):
|
|
summary = worker.execute()
|
|
|
|
self.assertEqual({"asset_id": 12, "cancelled": True}, summary)
|
|
|
|
def test_product_suite_ai_write_worker_uses_image_analysis_not_title_generation(self):
|
|
worker = ProductSuiteAiWriteWorker(
|
|
"补充要求",
|
|
"输出语言：繁体中文",
|
|
image_paths=["first.jpg", "second.jpg"],
|
|
config={"ai": {"backend": "cmhub"}},
|
|
cmhub_config_path="cmhub.json",
|
|
)
|
|
expected = {
|
|
"text": "根据图片整理的商品卖点",
|
|
"image_count": 2,
|
|
"metadata": {"points_cost": 1, "points_balance": 231},
|
|
}
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.ai.analyze_product_images",
|
|
return_value=expected,
|
|
) as analyze, mock.patch("app.gui.workers.ai.gen_title") as gen_title:
|
|
result = worker.execute()
|
|
|
|
analyze.assert_called_once_with(
|
|
"补充要求",
|
|
"输出语言：繁体中文",
|
|
["first.jpg", "second.jpg"],
|
|
config={"ai": {"backend": "cmhub"}},
|
|
cmhub_config_path="cmhub.json",
|
|
)
|
|
gen_title.assert_not_called()
|
|
self.assertEqual(expected, result)
|
|
|
|
def test_product_suite_history_export_worker_uses_round_export_service(self):
|
|
worker = ProductSuiteHistoryExportWorker(
|
|
7,
|
|
"round-key",
|
|
"D:/exports",
|
|
db_path="suite.db",
|
|
)
|
|
expected = SimpleNamespace(
|
|
target_dir="D:/exports/main-shop_51100639510_20260717",
|
|
files=(object(), object()),
|
|
skipped_count=1,
|
|
cancelled=False,
|
|
)
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio_export.export_generation_round",
|
|
return_value=expected,
|
|
) as export_round:
|
|
result = worker.execute()
|
|
|
|
export_round.assert_called_once_with(
|
|
7,
|
|
"round-key",
|
|
"D:/exports",
|
|
path="suite.db",
|
|
should_stop=worker.should_cancel,
|
|
)
|
|
self.assertEqual(
|
|
{
|
|
"target_dir": expected.target_dir,
|
|
"file_count": 2,
|
|
"skipped_count": 1,
|
|
"cancelled": False,
|
|
},
|
|
result,
|
|
)
|
|
|
|
def test_product_suite_worker_writes_generation_round_and_stable_slots(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = os.path.join(temp_dir, "cmshopee.db")
|
|
db.init_db(db_path)
|
|
project = image_studio.create_or_get_project(
|
|
account_alias="alias-a",
|
|
account_slug="alias-a",
|
|
item_id="51100639510",
|
|
path=db_path,
|
|
)
|
|
source = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
worker = ProductSuiteGenerateWorker(
|
|
project.id,
|
|
[
|
|
{
|
|
"source_asset_id": source.id,
|
|
"job_type": "白底图",
|
|
"prompt": "第一张",
|
|
"generation_round_key": "worker-round",
|
|
"generation_slot_index": 0,
|
|
},
|
|
{
|
|
"source_asset_id": source.id,
|
|
"job_type": "场景图",
|
|
"prompt": "第二张",
|
|
"generation_round_key": "worker-round",
|
|
"generation_slot_index": 1,
|
|
},
|
|
],
|
|
generation_round_key="worker-round",
|
|
db_path=db_path,
|
|
)
|
|
|
|
with mock.patch(
|
|
"app.gui.workers.image_studio_generation.run_jobs",
|
|
return_value={"total": 2, "success": 0, "failed": 0, "cancelled": 0},
|
|
):
|
|
summary = worker.execute()
|
|
|
|
jobs = [image_studio.get_job(job_id, path=db_path) for job_id in worker.job_ids]
|
|
self.assertEqual("worker-round", summary["generation_round_key"])
|
|
self.assertEqual(
|
|
[("worker-round", 0), ("worker-round", 1)],
|
|
[
|
|
(job.generation_round_key, job.generation_slot_index)
|
|
for job in jobs
|
|
],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|