cmshoppe/tests/test_workers.py

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 cmhub_models, 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 (
    CMHubModelCatalogWorker,
    ImageStudioDownloadOriginalWorker,
    ImageStudioPullImagesWorker,
    ProductSuiteAiWriteWorker,
    ProductSuiteGenerateWorker,
    ProductSuiteHistoryExportWorker,
    ProductTabOpenWorker,
    CollectWorker,
)


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_product_tab_open_worker_focuses_product_without_task_write(self):
        account = SimpleNamespace(
            alias="alias-a", account_name="主店", debug_port=9222
        )
        with mock.patch(
            "app.gui.workers.db.get_account_by_alias", return_value=account
        ) as get_account, mock.patch(
            "app.gui.workers.chrome.is_running", return_value=True
        ) as is_running, mock.patch(
            "app.gui.workers.editor.open_or_focus_product_tab",
            return_value={"created": False, "target_id": "target-existing"},
        ) as open_tab, mock.patch("app.gui.workers.db.set_generated") as set_generated:
            result = ProductTabOpenWorker(
                "alias-a", "51100639510", db_path="test.db"
            ).execute()

        self.assertTrue(result["ok"])
        self.assertFalse(result["created"])
        self.assertEqual("主店", result["account_name"])
        get_account.assert_called_once_with("alias-a", path="test.db")
        is_running.assert_called_once_with(9222)
        open_tab.assert_called_once_with(account, "51100639510")
        set_generated.assert_not_called()

    def test_product_tab_open_worker_reports_missing_chrome_without_starting_it(self):
        account = SimpleNamespace(
            alias="alias-a", account_name="主店", debug_port=9222
        )
        with mock.patch(
            "app.gui.workers.db.get_account_by_alias", return_value=account
        ), mock.patch(
            "app.gui.workers.chrome.is_running", return_value=False
        ), mock.patch("app.gui.workers.editor.open_or_focus_product_tab") as open_tab:
            result = ProductTabOpenWorker("alias-a", "51100639510").execute()

        self.assertFalse(result["ok"])
        self.assertEqual("CHROME_NOT_RUNNING", result["reason"])
        self.assertIn("④账号管理", result["message"])
        open_tab.assert_not_called()

    def test_collect_worker_scope_skips_non_normal_without_overwriting_old_content(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            batch_id = db.create_batch(["input.xlsx"], path=db_path)
            db.insert_tasks(
                batch_id,
                [
                    {
                        "source_file_abs": os.path.join(temp_dir, "input.xlsx"),
                        "source_sheet": "Sheet1",
                        "source_row": index,
                        "account_name": "店铺",
                        "alias": "alias",
                        "item_id": str(51100639510 + index),
                    }
                    for index in range(2, 6)
                ],
                path=db_path,
            )
            tasks = db.list_tasks(batch_id=batch_id, path=db_path)
            connection = db.connect(db_path)
            try:
                with connection:
                    connection.execute(
                        "UPDATE tasks SET old_title = ?, old_cover_path = ? WHERE id = ?",
                        ("历史标题", "history.jpg", tasks[1].id),
                    )
            finally:
                connection.close()

            status_by_item = {
                tasks[0].item_id: "normal",
                tasks[1].item_id: "unlisted",
                tasks[2].item_id: "reviewing",
                tasks[3].item_id: "unknown",
            }
            account = SimpleNamespace(alias="alias", account_name="店铺", debug_port=9222)

            def fake_collect(_account, task, on_step=None):
                status = status_by_item[task["item_id"]]
                self.assertEqual("normal_only", task["collection_scope"])
                on_step("read_product_status")
                if status != "normal":
                    labels = {
                        "unlisted": "未上架",
                        "reviewing": "审核中",
                        "unknown": "状态未知",
                    }
                    return {
                        "product_status": status,
                        "product_status_note": f"{labels[status]}提示",
                        "collection_skipped": True,
                        "collection_skip_reason": f"{labels[status]}，按本轮范围略过",
                    }
                on_step("download_cover")
                return {
                    "product_status": status,
                    "product_status_note": None,
                    "old_title": "新采集标题",
                    "old_cover_path": "new.jpg",
                }

            with mock.patch(
                "app.gui.workers.accounts.list_accounts",
                return_value=[account],
            ), mock.patch(
                "app.gui.workers.accounts.detect_login",
                return_value={"logged_in": True, "reason": None},
            ), mock.patch(
                "app.gui.workers.editor.collect",
                side_effect=fake_collect,
            ):
                summary = CollectWorker(
                    tasks,
                    db_path=db_path,
                    preflight=False,
                    collect_scope="normal_only",
                ).execute()

            self.assertEqual(1, summary["collected"])
            self.assertEqual(3, summary["skipped"])
            self.assertEqual(0, summary["failed"])
            self.assertEqual(3, summary["status_scope_skipped"])
            self.assertEqual(
                {"normal": 1, "unlisted": 1, "reviewing": 1, "unknown": 1},
                summary["product_status_counts"],
            )
            refreshed = db.list_tasks(batch_id=batch_id, path=db_path)
            by_status = {task.product_status: task for task in refreshed}
            self.assertEqual("collected", by_status["normal"].stage)
            self.assertEqual("success", by_status["normal"].status)
            self.assertEqual("skipped", by_status["unlisted"].status)
            self.assertEqual("imported", by_status["unlisted"].stage)
            self.assertEqual("历史标题", by_status["unlisted"].old_title)
            self.assertEqual("history.jpg", by_status["unlisted"].old_cover_path)

    def test_collect_worker_all_scope_collects_every_detected_status(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            batch_id = db.create_batch(["input.xlsx"], path=db_path)
            db.insert_tasks(
                batch_id,
                [
                    {
                        "source_file_abs": os.path.join(temp_dir, "input.xlsx"),
                        "source_sheet": "Sheet1",
                        "source_row": index,
                        "account_name": "店铺",
                        "alias": "alias",
                        "item_id": str(51100639600 + index),
                    }
                    for index in range(2, 6)
                ],
                path=db_path,
            )
            tasks = db.list_tasks(batch_id=batch_id, path=db_path)
            statuses = ["normal", "unlisted", "reviewing", "unknown"]
            account = SimpleNamespace(alias="alias", account_name="店铺", debug_port=9222)

            def fake_collect(_account, task, on_step=None):
                self.assertEqual("all", task["collection_scope"])
                on_step("read_product_status")
                status = statuses.pop(0)
                return {
                    "product_status": status,
                    "product_status_note": None,
                    "old_title": f"标题{status}",
                    "old_cover_path": f"{status}.jpg",
                }

            with mock.patch(
                "app.gui.workers.accounts.list_accounts",
                return_value=[account],
            ), mock.patch(
                "app.gui.workers.accounts.detect_login",
                return_value={"logged_in": True, "reason": None},
            ), mock.patch(
                "app.gui.workers.editor.collect",
                side_effect=fake_collect,
            ):
                summary = CollectWorker(
                    tasks,
                    db_path=db_path,
                    preflight=False,
                    collect_scope="all",
                ).execute()

            self.assertEqual(4, summary["collected"])
            self.assertEqual(0, summary["skipped"])
            self.assertEqual(0, summary["status_scope_skipped"])
            self.assertEqual(
                {"normal": 1, "unlisted": 1, "reviewing": 1, "unknown": 1},
                summary["product_status_counts"],
            )
            self.assertTrue(
                all(task.stage == "collected" for task in db.list_tasks(batch_id=batch_id, path=db_path))
            )

    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": "direct"}},
            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()
        args, kwargs = analyze.call_args
        self.assertEqual(
            ("补充要求", "输出语言：繁体中文", ["first.jpg", "second.jpg"]),
            args,
        )
        self.assertEqual(worker.config, kwargs["config"])
        self.assertEqual("cmhub.json", kwargs["cmhub_config_path"])
        self.assertIn("cmhub_api_key", worker.config["_cmshopee_ai_runtime"])
        self.assertNotIn("direct_models", worker.config["_cmshopee_ai_runtime"])
        gen_title.assert_not_called()
        self.assertEqual(expected, result)

    def test_cmhub_model_catalog_worker_caches_models_without_exposing_key(self):
        cmhub_models.clear_model_catalog_cache()
        self.addCleanup(cmhub_models.clear_model_catalog_cache)
        worker = CMHubModelCatalogWorker(
            "https://cmhub.example.com/",
            "sk-cmhub-secret",
            connect_timeout=7,
            use_system_proxy=True,
        )
        models = [{"alias": "vision-standard", "operation_type": "vision"}]

        with mock.patch("app.gui.workers.ai.fetch_cmhub_models", return_value=models) as fetch:
            result = worker.execute()

        fetch.assert_called_once_with(
            "https://cmhub.example.com",
            "sk-cmhub-secret",
            connect_timeout=7,
            use_system_proxy=True,
        )
        self.assertEqual(models, result["models"])
        self.assertEqual(
            models,
            cmhub_models.cached_model_catalog(
                "https://cmhub.example.com",
                "vision-standard",
            ),
        )

    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()