chengma/cmshoppe
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(product-suite): make image pulls cancellable

Browse Source
This commit is contained in:
chengma
2026-07-16 19:37:22 +08:00
parent c9afaad24b
commit 75af87d441
12 changed files with 1157 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/gui/tabs/product_suite.py
+459 -29
View File
@@ -707,6 +707,15 @@ class SuiteTaskState:
generation_terminal_streak: int = 0
pull_worker: object = None
pull_thread: object = None
pull_run_token: str = ""
pull_stop_requested: bool = False
pull_cleanup_mode: str = "keep"
pull_before_asset_ids: set = field(default_factory=set)
pull_before_asset_states: list = field(default_factory=list)
pull_asset_ids: set = field(default_factory=set)
pull_download_asset_ids: set = field(default_factory=set)
pull_download_failed: int = 0
pull_started_at: float = None
import_worker: object = None
import_thread: object = None
import_created_draft: bool = False
@@ -714,6 +723,7 @@ class SuiteTaskState:
ai_thread: object = None
download_queue: list = field(default_factory=list)
downloads: dict = field(default_factory=dict)
download_tokens: dict = field(default_factory=dict)
done: int = 0
failed: int = 0
total: int = 0
@@ -725,6 +735,9 @@ class SuiteTaskState:
def generation_running(self):
return self.worker is not None
def pull_running(self):
return bool(self.pull_run_token)
class ProductSuiteTab(QWidget):
"""Native PySide6 product-suite UI backed by image_studio services."""
@@ -757,6 +770,7 @@ class ProductSuiteTab(QWidget):
self._displayed_state = None
self._prompt_save_timers = {}
self._generation_run_states = {}
self._pull_run_states = {}
self._original_list_context = None
self._loading = False
self._result_refresh_pending = False
@@ -1363,8 +1377,12 @@ class ProductSuiteTab(QWidget):
state.worker.cancel()
if state.ai_worker is not None:
state.ai_worker.cancel()
if state.pull_running():
state.pull_stop_requested = True
state.pull_cleanup_mode = "keep"
if state.pull_worker is not None:
state.pull_worker.cancel()
self._cancel_pull_downloads(state)
if image_studio.is_draft_project(project):
if draft_action == "delete":
try:
@@ -1385,6 +1403,8 @@ class ProductSuiteTab(QWidget):
self._release_prompt_save_timer(state)
if state.generation_run_token:
self._generation_run_states.pop(state.generation_run_token, None)
if state.pull_run_token:
self._pull_run_states.pop(state.pull_run_token, None)
self._retired_states.append(state)
self._states.pop(state.key, None)
self.task_tabs.removeTab(index)
@@ -1602,6 +1622,13 @@ class ProductSuiteTab(QWidget):
state.generation_job_ids = []
state.generation_mode = "batch"
state.generation_retry_job_id = None
state.pull_run_token = ""
state.pull_stop_requested = False
state.pull_before_asset_ids = set()
state.pull_before_asset_states = []
state.pull_asset_ids = set()
state.pull_download_asset_ids = set()
state.pull_download_failed = 0
state.done = state.failed = state.total = 0
state.started_at = None
if state is self._displayed_state:
@@ -1916,13 +1943,23 @@ class ProductSuiteTab(QWidget):
def _update_context_actions(self, state):
is_draft = self._is_draft_state(state)
pull_running = bool(state and state.pull_running())
self.pull_button.setEnabled(
state is not None
and state.pull_worker is None
and (is_draft or self._valid_context(state, show_message=False))
and (
pull_running
or is_draft
or self._valid_context(state, show_message=False)
)
)
self.pull_button.setToolTip(
"需要先绑定正式商品ID" if is_draft else "拉取蝦皮主图"
"正在停止拉取"
if pull_running and state.pull_stop_requested
else "停止拉取蝦皮主图"
if pull_running
else "需要先绑定正式商品ID"
if is_draft
else "拉取蝦皮主图"
)
self._refresh_item_id_hint(state)
self._refresh_add_images_action(state)
@@ -1982,6 +2019,10 @@ class ProductSuiteTab(QWidget):
if run_token:
thread.setProperty("productSuiteRunToken", run_token)
thread.finished.connect(self._on_generation_thread_finished_signal)
pull_run_token = str(getattr(worker, "pull_run_token", "") or "")
if pull_run_token:
thread.setProperty("productSuitePullRunToken", pull_run_token)
thread.finished.connect(self._on_pull_thread_finished_signal)
thread.start()
return thread
@@ -2158,6 +2199,9 @@ class ProductSuiteTab(QWidget):
if state.generation_running():
self._status("生成中不能删除当前任务的商品原图", "warning")
return
if state.pull_running():
self._status("拉取蝦皮主图时不能删除当前任务的商品原图", "warning")
return
normalized_ids = []
seen = set()
for value in asset_ids or []:
@@ -2243,6 +2287,9 @@ class ProductSuiteTab(QWidget):
state = self._displayed_state
if state is None:
return
if state.pull_running():
self._request_stop_pull(state)
return
self._save_controls_to_state(state)
if self._is_draft_state(state):
self._message(
@@ -2252,9 +2299,6 @@ class ProductSuiteTab(QWidget):
return
if not self._valid_context(state):
return
if state.pull_worker is not None:
self._status("当前任务正在拉取蝦皮主图", "info")
return
existing = [asset for asset in self._original_assets(state, include_missing=False) if _asset_usable(asset)]
message = (
"店铺：%s\n"
@@ -2276,33 +2320,172 @@ class ProductSuiteTab(QWidget):
cancel_text="取消",
):
return
if state.project_id is None:
project = image_studio.get_project_by_account_item(
state.account_alias,
state.item_id,
path=self.db_path,
)
if project is not None:
state.project_id = int(project.id)
state.project_binding_state = project.binding_state
before_assets = (
image_studio.list_assets(
state.project_id,
kind=image_studio.ASSET_KIND_ORIGINAL,
path=self.db_path,
)
if state.project_id is not None
else []
)
pull_run_token = uuid.uuid4().hex
worker = ImageStudioPullImagesWorker(
state.account_alias,
state.item_id,
pull_run_token=pull_run_token,
db_path=self.db_path,
config=self.config,
)
state.pull_worker = worker
worker.finished.connect(lambda result, state=state: self._on_pull_finished(state, result))
worker.cancelled.connect(lambda result, state=state: self._on_pull_finished(state, result))
worker.failed.connect(lambda row, error, state=state: self._on_pull_failed(state, error))
state.pull_run_token = pull_run_token
state.pull_stop_requested = False
state.pull_cleanup_mode = "keep"
state.pull_before_asset_ids = {int(asset.id) for asset in before_assets}
state.pull_before_asset_states = [
{
"id": int(asset.id),
"status": asset.status,
"source_order": int(asset.source_order or 0),
}
for asset in before_assets
]
state.pull_asset_ids = set()
state.pull_download_asset_ids = set()
state.pull_download_failed = 0
state.pull_started_at = time.monotonic()
self._pull_run_states[pull_run_token] = state.key
self._log_pull_lifecycle(
state,
pull_run_token,
"started",
{
"before_count": len(state.pull_before_asset_ids),
},
)
worker.finished.connect(
lambda result, token=pull_run_token: self._on_pull_finished(
token,
result,
)
)
worker.cancelled.connect(
lambda result, token=pull_run_token: self._on_pull_finished(
token,
result,
)
)
worker.failed.connect(
lambda row, error, token=pull_run_token: self._on_pull_failed(
token,
error,
)
)
state.pull_thread = self._start_thread(worker, "商品套图拉取蝦皮主图")
if state is self._displayed_state:
self.pull_button.setText("正在拉取...")
self._update_context_actions(state)
self._apply_running_state(state)
self._status("开始拉取蝦皮主图，可继续操作其他套图任务", "info")
def _on_pull_failed(self, state, error):
def _pull_state(self, pull_run_token):
token = str(pull_run_token or "")
state = self._states.get(self._pull_run_states.get(token))
if state is None or state.pull_run_token != token:
return None
return state
def _pull_stop_action(self):
box = QMessageBox(self)
box.setIcon(QMessageBox.Question)
box.setWindowTitle("停止拉取蝦皮主图")
box.setText("请选择停止后如何处理本轮已经拉取的图片。")
keep_button = box.addButton("停止并保留", QMessageBox.AcceptRole)
clear_button = box.addButton(
"停止并清除本次新增",
QMessageBox.DestructiveRole,
)
continue_button = box.addButton("继续拉取", QMessageBox.RejectRole)
box.setDefaultButton(keep_button)
box.exec()
clicked = box.clickedButton()
if clicked is keep_button:
return "keep"
if clicked is clear_button:
return "clear_current"
if clicked is continue_button:
return "continue"
return "continue"
def _request_stop_pull(self, state):
if state.pull_stop_requested:
self._status("正在停止当前拉取任务", "warning")
return
action = self._pull_stop_action()
if action == "continue":
return
state.pull_stop_requested = True
state.pull_cleanup_mode = action
if state.pull_worker is not None:
state.pull_worker.cancel()
self._cancel_pull_downloads(state)
self._log_pull_lifecycle(
state,
state.pull_run_token,
"stop_requested",
{"cleanup_mode": action},
level="WARNING",
)
if state is self._displayed_state:
self._apply_running_state(state)
self._status("已请求停止拉取蝦皮主图", "warning")
self._maybe_finalize_pull(state, state.pull_run_token)
def _cancel_pull_downloads(self, state):
pull_ids = set(state.pull_download_asset_ids)
if not pull_ids:
return
state.download_queue = [
asset_id
for asset_id in state.download_queue
if int(asset_id) not in pull_ids
]
for asset_id in list(pull_ids):
running = state.downloads.get(int(asset_id))
if running is None:
state.pull_download_asset_ids.discard(int(asset_id))
state.download_tokens.pop(int(asset_id), None)
continue
worker, _ = running
worker.cancel()
def _on_pull_failed(self, pull_run_token, error):
state = self._pull_state(pull_run_token)
if state is None or state.pull_stop_requested:
return
self._status("拉取蝦皮主图失败：%s" % _user_error(error), "danger")
def _on_pull_finished(self, state, result):
def _on_pull_finished(self, pull_run_token, result):
state = self._pull_state(pull_run_token)
if state is None:
return
state.pull_worker = None
state.pull_thread = None
if state is self._displayed_state:
self.pull_button.setText("拉取蝦皮主图")
if result.get("ok") is False:
self._message("拉取蝦皮主图失败", _user_error(result.get("error")))
else:
result = dict(result or {})
if result.get("ok") is False and not state.pull_stop_requested:
self._finalize_pull(
state,
pull_run_token,
error=_user_error(result.get("error")),
)
return
project = result.get("project")
if project is not None:
state.project_id = int(project.id)
@@ -2312,13 +2495,231 @@ class ProductSuiteTab(QWidget):
for asset in (result.get("assets") or [])
if asset.status != image_studio.ASSET_STATUS_MISSING
][: image_studio_images.MAX_ORIGINAL_ASSETS]
self._queue_original_downloads(state, assets)
state.pull_asset_ids.update(
int(asset.id)
for asset in assets
if getattr(asset, "remote_url", None)
)
if result.get("cancelled") or state.pull_stop_requested:
self._cancel_pull_downloads(state)
else:
self._queue_original_downloads(
state,
assets,
pull_run_token=pull_run_token,
)
self._status("已读取%d张蝦皮主图，正在后台下载" % len(assets), "success")
if state is self._displayed_state:
self._refresh_originals(state)
self._update_context_actions(state)
self._apply_running_state(state)
self._maybe_finalize_pull(state, pull_run_token)
def _queue_original_downloads(self, state, assets):
@Slot()
def _on_pull_thread_finished_signal(self):
sender = self.sender()
token = str(
sender.property("productSuitePullRunToken")
if sender is not None
else ""
)
QTimer.singleShot(0, lambda token=token: self._handle_pull_thread_finished(token))
def _handle_pull_thread_finished(self, pull_run_token):
state = self._pull_state(pull_run_token)
if state is None or state.pull_worker is None:
return
state.pull_worker = None
state.pull_thread = None
if state.pull_stop_requested:
self._maybe_finalize_pull(state, pull_run_token)
else:
self._finalize_pull(
state,
pull_run_token,
error="拉取线程已结束，请稍后重试",
)
def _maybe_finalize_pull(self, state, pull_run_token):
if self._pull_state(pull_run_token) is not state:
return False
if state.pull_worker is not None or state.pull_download_asset_ids:
return False
return self._finalize_pull(state, pull_run_token)
def _clear_current_pull_assets(self, state):
new_ids = sorted(
set(state.pull_asset_ids) - set(state.pull_before_asset_ids)
)
removable = []
retained = 0
for asset_id in new_ids:
asset = image_studio.get_asset(asset_id, path=self.db_path)
if (
asset is None
or int(asset.project_id) != int(state.project_id or 0)
or asset.kind != image_studio.ASSET_KIND_ORIGINAL
or not asset.remote_url
):
continue
if image_studio.asset_reference_counts(
asset.id,
path=self.db_path,
)["total"]:
retained += 1
else:
removable.append(asset.id)
removed = 0
if removable:
image_studio.remove_original_assets_if_unused(
state.project_id,
removable,
path=self.db_path,
)
removed = len(removable)
existing_ids = {
int(asset.id)
for asset in image_studio.list_assets(
state.project_id,
kind=image_studio.ASSET_KIND_ORIGINAL,
path=self.db_path,
)
}
snapshot = [
item
for item in state.pull_before_asset_states
if int(item["id"]) in existing_ids
]
if snapshot:
image_studio.restore_original_asset_snapshot(
state.project_id,
snapshot,
path=self.db_path,
)
return removed, retained
def _finalize_pull(self, state, pull_run_token, *, error=""):
if self._pull_state(pull_run_token) is not state:
return False
cancelled = state.pull_stop_requested
cleanup_mode = state.pull_cleanup_mode
pulled_count = len(state.pull_asset_ids)
new_count = len(set(state.pull_asset_ids) - set(state.pull_before_asset_ids))
removed = 0
retained = 0
cleanup_error = ""
if cancelled and cleanup_mode == "clear_current" and state.project_id is not None:
try:
removed, retained = self._clear_current_pull_assets(state)
except Exception as exc:
cleanup_error = _user_error(exc)
elapsed = int(
max(0, time.monotonic() - state.pull_started_at)
if state.pull_started_at
else 0
)
self._pull_run_states.pop(pull_run_token, None)
state.pull_worker = None
state.pull_thread = None
state.pull_run_token = ""
state.pull_stop_requested = False
state.pull_cleanup_mode = "keep"
state.pull_before_asset_ids = set()
state.pull_before_asset_states = []
state.pull_asset_ids = set()
state.pull_download_asset_ids = set()
failed = int(state.pull_download_failed or 0)
self._log_pull_lifecycle(
state,
pull_run_token,
"finalized",
{
"cancelled": cancelled,
"cleanup_mode": cleanup_mode,
"pulled_count": pulled_count,
"new_count": new_count,
"removed_count": removed,
"retained_count": retained,
"failed_count": failed,
"elapsed_seconds": elapsed,
"has_error": bool(error or cleanup_error),
},
level="WARNING" if cancelled or error or cleanup_error else "INFO",
)
state.pull_download_failed = 0
state.pull_started_at = None
if state is self._displayed_state:
self._refresh_originals(state)
self._apply_running_state(state)
if cleanup_error:
self._message(
"停止拉取后清理失败",
"拉取已停止，但本轮图片清理失败：%s" % cleanup_error,
)
self._status("拉取已停止，本轮图片清理失败", "danger")
return True
if error:
if state is self._displayed_state:
self._message("拉取蝦皮主图失败", error)
self._status("拉取蝦皮主图失败：%s" % error, "danger")
return True
if cancelled:
kept = max(0, new_count - removed)
if state is self._displayed_state:
self._message(
"拉取蝦皮主图已停止",
"本轮新增%d张：保留%d张，清理%d张，"
"因引用保留%d张，下载失败%d张；总用时%d秒。"
% (
new_count,
kept,
removed,
retained,
failed,
elapsed,
),
icon=QMessageBox.Information,
)
self._status(
"拉取已停止：保留%d张，清理%d张" % (kept, removed),
"warning",
)
return True
self._status(
"蝦皮主图拉取完成：本轮%d张，下载失败%d张，用时%d秒"
% (pulled_count, failed, elapsed),
"success" if not failed else "warning",
)
return True
def _log_pull_lifecycle(
self,
state,
pull_run_token,
event,
payload=None,
*,
level="INFO",
):
data = {
"pull_run_token": str(pull_run_token or "")[:8],
"project_id": getattr(state, "project_id", None),
"event": str(event or ""),
}
data.update(dict(payload or {}))
try:
diagnostics.write_diagnostic_log(
"商品套图拉取生命周期",
level=level,
step="product_suite_pull",
task_id=getattr(state, "project_id", None),
item_id=getattr(state, "item_id", None),
payload=data,
log_dir=appconfig.diagnostic_log_dir(self.config),
)
except Exception:
pass
def _queue_original_downloads(self, state, assets, *, pull_run_token=None):
if state is None:
return
queued = set(state.download_queue) | set(state.downloads)
@@ -2328,6 +2729,9 @@ class ProductSuiteTab(QWidget):
asset_id = int(asset.id)
if asset_id not in queued:
state.download_queue.append(asset_id)
state.download_tokens[asset_id] = str(pull_run_token or "")
if pull_run_token:
state.pull_download_asset_ids.add(asset_id)
queued.add(asset_id)
self._start_queued_downloads(state)
@@ -2340,14 +2744,15 @@ class ProductSuiteTab(QWidget):
config=self.config,
max_retries=2,
)
download_token = str(state.download_tokens.get(asset_id, "") or "")
worker.finished.connect(
lambda result, state=state, asset_id=asset_id: self._on_download_finished(
state, asset_id, result
lambda result, state=state, asset_id=asset_id, token=download_token: self._on_download_finished(
state, asset_id, result, pull_run_token=token
)
)
worker.cancelled.connect(
lambda result, state=state, asset_id=asset_id: self._on_download_finished(
state, asset_id, result
lambda result, state=state, asset_id=asset_id, token=download_token: self._on_download_finished(
state, asset_id, result, pull_run_token=token
)
)
worker.failed.connect(
@@ -2359,9 +2764,24 @@ class ProductSuiteTab(QWidget):
thread = self._start_thread(worker, "商品套图下载原图")
state.downloads[asset_id] = (worker, thread)
def _on_download_finished(self, state, asset_id, result):
def _on_download_finished(
self,
state,
asset_id,
result,
*,
pull_run_token="",
):
expected_token = str(state.download_tokens.get(int(asset_id), "") or "")
if expected_token != str(pull_run_token or ""):
return
state.downloads.pop(int(asset_id), None)
state.download_tokens.pop(int(asset_id), None)
if pull_run_token:
state.pull_download_asset_ids.discard(int(asset_id))
if result.get("ok") is False:
if pull_run_token:
state.pull_download_failed += 1
self._status(
"商品原图 #%d 下载失败：%s" % (asset_id, _user_error(result.get("error"))),
"danger",
@@ -2371,6 +2791,8 @@ class ProductSuiteTab(QWidget):
if state is self._displayed_state:
self._refresh_originals(state)
self._start_queued_downloads(state)
if pull_run_token:
self._maybe_finalize_pull(state, pull_run_token)
def _rebuild_categories(self, state):
while self.category_grid.count():
@@ -3221,9 +3643,16 @@ class ProductSuiteTab(QWidget):
def _apply_running_state(self, state):
generation_running = state.generation_running()
self.pull_button.setText("正在拉取..." if state.pull_worker is not None else "拉取蝦皮主图")
self.account_combo.setEnabled(not generation_running and state.pull_worker is None)
self.item_id_edit.setEnabled(not generation_running and state.pull_worker is None)
pull_running = state.pull_running()
self.pull_button.setText(
"正在停止..."
if pull_running and state.pull_stop_requested
else "停止拉取蝦皮"
if pull_running
else "拉取蝦皮主图"
)
self.account_combo.setEnabled(not generation_running and not pull_running)
self.item_id_edit.setEnabled(not generation_running and not pull_running)
self._refresh_add_images_action(state)
self.original_list.setEnabled(not generation_running)
self._refresh_original_selection_controls()
@@ -3522,4 +3951,5 @@ class ProductSuiteTab(QWidget):
for state in list(self._states.values()) + list(self._retired_states):
self._release_prompt_save_timer(state)
self._generation_run_states.clear()
self._pull_run_states.clear()
super().closeEvent(event)
app/gui/workers.py
+31 -1
View File
@@ -90,31 +90,58 @@ def _image_studio_user_detail(detail):
class ImageStudioPullImagesWorker(BaseWorker):
"""Read Shopee main image URLs for one AI studio project in background."""
def __init__(self, account_alias, item_id, *, db_path=None, config=None):
def __init__(
self,
account_alias,
item_id,
*,
pull_run_token="",
db_path=None,
config=None,
):
super().__init__()
self.account_alias = account_alias
self.item_id = item_id
self.pull_run_token = str(pull_run_token or "")
self.db_path = db_path
self.config = config
def execute(self):
if self.should_cancel():
return {
"pull_run_token": self.pull_run_token,
"cancelled": True,
"assets": [],
}
def on_step(payload):
self.log.emit(_format_image_studio_event(payload))
try:
result = image_studio.pull_remote_main_image_urls(
self.account_alias,
self.item_id,
path=self.db_path,
config=self.config,
on_step=on_step,
should_stop=self.should_cancel,
)
except image_studio.ImageStudioPullCancelled as exc:
return {
"pull_run_token": self.pull_run_token,
"cancelled": True,
"project": exc.project,
"assets": list(exc.assets or []),
}
project = result.get("project")
assets = result.get("assets") or []
return {
"pull_run_token": self.pull_run_token,
"project": project,
"assets": assets,
"count": len(assets),
"account": result.get("account"),
"cancelled": self.should_cancel(),
}
@@ -158,7 +185,10 @@ class ImageStudioDownloadOriginalWorker(BaseWorker):
self.asset_id,
path=self.db_path,
config=self.config,
should_stop=self.should_cancel,
)
except image_studio_images.ImageStudioImageCancelled:
return {"asset_id": self.asset_id, "cancelled": True}
except Exception:
retry = attempt
if attempt >= attempts:
app/image_studio.py
+112
View File
@@ -116,6 +116,15 @@ class ImageStudioProjectConflictError(ImageStudioError):
"""Raised when a draft cannot be bound because the formal project already exists."""
class ImageStudioPullCancelled(ImageStudioError):
"""Raised when a read-only Shopee image pull stops at a safe boundary."""
def __init__(self, project=None, assets=None):
super().__init__("用户停止拉取蝦皮主图")
self.project = project
self.assets = list(assets or [])
def _now() -> str:
return datetime.now().isoformat(timespec="seconds")
@@ -825,6 +834,93 @@ def remove_original_assets_if_unused(project_id, asset_ids, path=None, conn=None
return [_row_to_dataclass(by_id[asset_id], ImageStudioAsset) for asset_id in ordered_ids]
def restore_original_asset_snapshot(project_id, snapshots, path=None, conn=None):
"""Restore status/order for original assets that existed before a pull."""
project_id = int(project_id)
normalized = []
seen = set()
for value in snapshots or []:
asset_id = int(_get(value, "id"))
if asset_id in seen:
continue
status = str(_get(value, "status") or ASSET_STATUS_AVAILABLE)
if status not in ASSET_STATUSES:
raise db.DbError("商品原图快照状态无效")
normalized.append(
{
"id": asset_id,
"status": status,
"source_order": max(0, int(_get(value, "source_order") or 0)),
}
)
seen.add(asset_id)
if not normalized:
return []
ids = [item["id"] for item in normalized]
placeholders = ",".join("?" for _ in ids)
with _connection(conn, path) as database:
with database:
rows = database.execute(
f"""
SELECT id FROM image_studio_assets
WHERE project_id = ? AND kind = ? AND id IN ({placeholders})
""",
[project_id, ASSET_KIND_ORIGINAL, *ids],
).fetchall()
if {int(row["id"]) for row in rows} != set(ids):
raise db.DbError("拉取前商品原图快照已失效")
now = _now()
for item in normalized:
database.execute(
"""
UPDATE image_studio_assets
SET status = ?, source_order = ?, updated_at = ?
WHERE id = ? AND project_id = ? AND kind = ?
""",
(
item["status"],
item["source_order"],
now,
item["id"],
project_id,
ASSET_KIND_ORIGINAL,
),
)
remaining_rows = database.execute(
f"""
SELECT id FROM image_studio_assets
WHERE project_id = ? AND kind = ?
AND id NOT IN ({placeholders})
ORDER BY source_order, id
""",
[project_id, ASSET_KIND_ORIGINAL, *ids],
).fetchall()
next_order = max(
[int(item["source_order"]) for item in normalized] + [0]
) + 1
for row in remaining_rows:
database.execute(
"""
UPDATE image_studio_assets
SET source_order = ?, updated_at = ?
WHERE id = ? AND project_id = ? AND kind = ?
""",
(
next_order,
now,
int(row["id"]),
project_id,
ASSET_KIND_ORIGINAL,
),
)
next_order += 1
return [
get_asset(item["id"], conn=database)
for item in normalized
]
def sync_original_asset_urls(project_id, image_urls, path=None, conn=None, max_assets=16):
"""Store the read-only Shopee main image URL snapshot as remote-only assets."""
@@ -1183,10 +1279,14 @@ def pull_remote_main_image_urls(
config=None,
login_timeout=8,
on_step=None,
should_stop=None,
):
"""Create/open an AI studio project and read Shopee main image URLs without editing."""
cfg = appconfig.load_config() if config is None else config
should_stop = should_stop or (lambda: False)
if should_stop():
raise ImageStudioPullCancelled()
database_path = _db_path(path, cfg)
db.init_db(database_path)
item = _normalize_item_id(item_id)
@@ -1196,6 +1296,8 @@ def pull_remote_main_image_urls(
raise ImageStudioError(f"AI工场账号不可用：{exc}") from exc
project = create_or_get_project(account, item_id=item, path=database_path)
if should_stop():
raise ImageStudioPullCancelled(project=project)
readiness = _ensure_account_ready_for_read(
account,
path=database_path,
@@ -1203,16 +1305,26 @@ def pull_remote_main_image_urls(
login_timeout=login_timeout,
on_step=on_step,
)
if should_stop():
raise ImageStudioPullCancelled(project=project)
cdp = None
try:
if should_stop():
raise ImageStudioPullCancelled(project=project)
_notify_step(on_step, "open_product", "start", f"商品 {item}")
cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False)
if should_stop():
raise ImageStudioPullCancelled(project=project)
_notify_step(on_step, "open_product", "success", f"商品 {item}")
_notify_step(on_step, "read_main_images", "start", f"商品 {item}")
images = editor.read_product_image_urls(cdp)
if should_stop():
raise ImageStudioPullCancelled(project=project)
if not images:
raise ImageStudioError("未读取到蝦皮商品主图 URL")
assets = sync_original_asset_urls(project.id, images, path=database_path)
if should_stop():
raise ImageStudioPullCancelled(project=project, assets=assets)
_notify_step(on_step, "read_main_images", "success", f"读取 {len(images)} 张主图 URL")
project = get_project(project.id, path=database_path)
return {
app/image_studio_images.py
+51 -4
View File
@@ -32,6 +32,10 @@ class ImageStudioImageError(RuntimeError):
"""Raised when a remote image cannot be safely loaded or saved."""
class ImageStudioImageCancelled(ImageStudioImageError):
"""Raised when an original image download is cancelled safely."""
@dataclass(frozen=True)
class RemoteImage:
url: str
@@ -127,10 +131,14 @@ def download_remote_image(
max_bytes=ORIGINAL_MAX_BYTES,
timeout=(DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS),
session=None,
should_stop=None,
) -> RemoteImage:
"""Download remote image bytes with SSRF, timeout and size guards."""
url = str(url or "").strip()
should_stop = should_stop or (lambda: False)
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
_assert_public_http_url(url)
client = session or _session()
try:
@@ -142,12 +150,17 @@ def download_remote_image(
)
except requests.exceptions.RequestException as exc:
raise ImageStudioImageError(f"远程图片下载失败：{exc}") from exc
try:
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
status = getattr(response, "status_code", 200)
if status >= 400:
raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}")
final_url = str(getattr(response, "url", "") or url)
_assert_public_http_url(final_url)
content_type = _validate_content_type(getattr(response, "headers", {}).get("Content-Type", ""))
content_type = _validate_content_type(
getattr(response, "headers", {}).get("Content-Type", "")
)
content_length = getattr(response, "headers", {}).get("Content-Length")
if content_length:
try:
@@ -157,14 +170,22 @@ def download_remote_image(
pass
chunks = []
total = 0
iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
iterator = (
response.iter_content(chunk_size=65536)
if hasattr(response, "iter_content")
else [response.content]
)
for chunk in iterator:
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
if not chunk:
continue
total += len(chunk)
if total > int(max_bytes):
raise ImageStudioImageError("远程图片超过大小上限")
chunks.append(chunk)
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
content = b"".join(chunks)
_image_info(content)
return RemoteImage(
@@ -174,6 +195,10 @@ def download_remote_image(
final_url=final_url,
redirected=final_url != url,
)
finally:
close = getattr(response, "close", None)
if callable(close):
close()
def _image_info(image_bytes):
@@ -509,10 +534,21 @@ def restore_trashed_asset(record, *, path=None, config=None):
raise ImageStudioImageError(f"撤销删除生成图片失败：{exc}") from exc
def download_original_asset(asset_id, *, path=None, config=None, image_root=None, session=None):
def download_original_asset(
asset_id,
*,
path=None,
config=None,
image_root=None,
session=None,
should_stop=None,
):
"""Download one Shopee original image into originals/ and mark its asset available."""
cfg = appconfig.load_config() if config is None else config
should_stop = should_stop or (lambda: False)
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
database_path = path or appconfig.db_path(cfg)
asset = image_studio.get_asset(asset_id, path=database_path)
if asset is None:
@@ -525,15 +561,26 @@ def download_original_asset(asset_id, *, path=None, config=None, image_root=None
project = image_studio.get_project(asset.project_id, path=database_path)
if project is None:
raise ImageStudioImageError("原图资产所属项目不存在")
remote = download_remote_image(asset.remote_url, max_bytes=ORIGINAL_MAX_BYTES, session=session)
remote = download_remote_image(
asset.remote_url,
max_bytes=ORIGINAL_MAX_BYTES,
session=session,
should_stop=should_stop,
)
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
info = _image_info(remote.content)
directory, stem = _original_file_path(project, asset, image_root or appconfig.image_dir(cfg))
os.makedirs(directory, exist_ok=True)
final_path = os.path.join(directory, stem + _extension_for_format(info["format"]))
temp_path = final_path + ".tmp-" + uuid.uuid4().hex
try:
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
with open(temp_path, "wb") as fh:
fh.write(remote.content)
if should_stop():
raise ImageStudioImageCancelled("用户停止下载商品原图")
with open(temp_path, "rb") as fh:
_image_info(fh.read())
os.replace(temp_path, final_path)
docs/04-architecture.md
+1
View File
@@ -438,6 +438,7 @@ data/images/<batch_id>/<slug>/<task_id>_<item_id>_new.<ext> # AI 生成的新
- `image_studio_projects.suite_settings_json` 持久化套图设置，旧数据库由 `db.init_db()` 原位补列，默认 `{}`；`draft_prompt` 继续保存卖点文本。`image_studio_assets` 中有效商品原图最多16张，历史 missing 记录不占有效名额；手工原图不会因再次同步蝦皮 URL 被误标 missing。⑥原图列表的批量勾选只保存在当前 `SuiteTaskState` 对应的界面上下文，不写库；批量移除由 `remove_original_assets_if_unused()` 一次校验项目归属、原图类型和 job/终选引用，并在单个 SQLite 事务中删除资产行、连续重排 `source_order`。服务不删除本地文件或蝦皮线上图片，任一资产校验失败时整批回滚。
- T-636 起，`image_studio_projects` 增加 `binding_state`（`draft` / `bound`）和稳定 `storage_key`。既有项目迁移为 `bound`，并以原 `item_id` 回填 `storage_key`；项目目录改用 `storage_key`，因此临时草稿绑定正式商品 ID 后不移动目录、不改写已有资产路径。草稿内部使用 `draft_<uuid>` 作为仅数据库可见的非空 `item_id`，GUI 输入框始终保持空白，用户日志和 cmhub 提示词只使用“临时草稿”或“未绑定商品”，不得暴露该内部值。
- ⑥已选账号但未填写商品 ID 时允许输入卖点、导入、拖入或粘贴本地图片。非空卖点在现有防抖稳定后会创建一个可恢复临时草稿并保存到 `image_studio_projects.draft_prompt`；首次有效图片导入也会创建草稿。空白卖点、取消选择和全部导入失败不保留空草稿。只有卖点的草稿同样属于可恢复业务内容；卖点清空且没有资产/job 时仍可按空草稿规则清理。草稿可管理本地图片、AI 帮写、生成套图、查看历史和打开结果目录，但在创建 worker、启动 Chrome 或执行 CDP 前禁止「拉取蝦皮主图」。输入合法数字商品 ID 后，经确认原地绑定同一个 `project_id`；资产、job、selection、提示词、套图设置和 `storage_key` 均保持不变。若同账号目标 ID（含软删除项目）已存在则拒绝覆盖或合并。
- ⑥拉取蝦皮主图使用独立内存 `pull_run_token` 覆盖 URL 读取和本轮后台下载。运行中按钮提供「停止并保留 / 停止并清除本次新增 / 继续拉取」；停止采用协作式安全边界，不强杀线程、Chrome 或已发出的 CDP/HTTP 请求。选择清除时只移除本轮新增、属于当前项目且未被 job/selection 引用的远程原图，并恢复拉取前已有原图的状态和顺序；本地手动导入图片、拉取前已有图片、用户文件和蝦皮线上图片不删除。旧 token 的 URL/下载迟到结果不得覆盖新一轮状态。
- 启动时恢复未软删除、至少含一条资产或生成任务的草稿为独立中文“临时草稿”标签，按最近更新时间排序。关闭非空草稿可选择保留、软删除或取消；软删除不物理删除图片目录。③「更新蝦皮」只处理正式任务，不接受临时草稿。
- 第六 Tab 的多个 `SuiteTaskState` 各自保留 generation/pull/import/AI/download worker 与线程引用；切换任务不取消任务。多个任务可并行，但 `image_studio_generation` 使用进程级 semaphore 保证所有套图任务合计最多5个 cmhub 在途 job。线程还在运行时关闭任务只请求协作式停止，模块级引用保留到 `QThread.finished`，不得提前销毁线程对象；下载前后均检查停止信号，停止后的临时文件不入资产库。
- T-639 后每轮套图生成使用仅存在内存的 `run_token` 隔离迟到信号，progress/finished/cancelled/failed 通过主线程绑定槽统一处理；正常 worker 结果、`QThread.finished` 和本轮 job 连续两次全部终态看门狗共同进入幂等 finalize。GUI 只按本轮明确 `job_ids` 判断完成，不用历史图片数量；即使最终信号丢失也会恢复按钮，旧线程引用仍保留到真实结束。停止为协作式：调度循环约每200ms检查标记并取消未开始 future，提交/轮询在有界请求返回后停止；requests 在流式数据块边界取消，Windows curl 由隐藏窗口 `Popen` 有界 terminate/kill。已有 `task_id` 的停止任务保留 resume，不假设服务端任务被取消或点数退回。
docs/api.md
+5 -3
View File
@@ -378,12 +378,14 @@ create_job(project_id, source_asset_id=None, job_type="main", prompt="", ...) ->
list_jobs(project_id, statuses=None, path=None) -> list[ImageStudioJob]
list_resumable_jobs(project_id=None, include_failed_downloads=False, path=None) -> list[ImageStudioJob]
replace_selections(project_id, selection_type, asset_ids, path=None) -> list[ImageStudioSelection]
pull_remote_main_image_urls(account_or_alias, item_id, path=None, config=None) -> dict
pull_remote_main_image_urls(account_or_alias, item_id, path=None, config=None,
should_stop=None) -> dict
# app/image_studio_images.py
download_remote_image(url, max_bytes=..., timeout=(connect, read)) -> RemoteImage
load_thumbnail(url, key=None, max_size=220) -> ThumbnailResult
download_original_asset(asset_id, path=None, config=None) -> ImageStudioAsset
download_original_asset(asset_id, path=None, config=None,
should_stop=None) -> ImageStudioAsset
import_original_files(project_id, file_paths, path=None, config=None) -> dict
import_original_bytes(project_id, content, filename_hint="clipboard.png", ...) -> ImageStudioAsset
trash_generated_asset(asset_id, path=None, config=None) -> dict
@@ -433,7 +435,7 @@ class GenerateWorker(BaseWorker) # ② 后台生成：ai.generate_batc
class ApplyWorker(BaseWorker) # ③ 后台更新：账号就绪预检 -> 检查或按批调用 editor.apply_task(...) -> db.set_applied/mark_skipped
class WriteBackWorker(BaseWorker) # ①/③ 后台回写：旧字段或更新结果写回原 Excel
class AIModelTestWorker(BaseWorker) # ⑤ 后台测试 AI 模型连接：appconfig.test_ai_model
class ImageStudioPullImagesWorker(BaseWorker) # ⑥ 后台只读拉蝦皮原主图 URL
class ImageStudioPullImagesWorker(BaseWorker) # ⑥ 后台只读拉蝦皮原主图 URL；安全边界协作停止
class ImageStudioDownloadOriginalWorker(BaseWorker)# ⑥ 后台下载远程原图
class ImageStudioGenerateJobsWorker(BaseWorker) # ⑥ 后台提交/查询/下载 cmhub 生图 job
class ImageStudioResumeJobsWorker(BaseWorker) # ⑥ 后台恢复已有 task_id 的生图 job
docs/routes.md
+1 -1
View File
@@ -243,6 +243,6 @@
| `ApplyWorker(BaseWorker)` | ③ | 账号就绪预检、检查本轮更新、按每批最大条数分批、按账号并行或串行调用 `editor.apply_task(...)`、逐条 `set_applied()`，失败继续，写运行日志 |
| `AIModelTestWorker(BaseWorker)` | ⑤ | 后台调用 `appconfig.test_ai_model()` 测试模型连接 |
| `WriteBackWorker(BaseWorker)` | ①③ | ①回写旧字段；③回写新标题/新封面/更新状态 |
| `ImageStudioPullImagesWorker / ImageStudioDownloadOriginalWorker / ProductSuiteImportImagesWorker / ProductSuiteGenerateWorker / ProductSuiteAiWriteWorker` | ⑥ | 后台执行只读拉主图、远程原图下载、本地图片校验复制、cmhub 套图生成与AI帮写；不直接操作 QWidget |
| `ImageStudioPullImagesWorker / ImageStudioDownloadOriginalWorker / ProductSuiteImportImagesWorker / ProductSuiteGenerateWorker / ProductSuiteAiWriteWorker` | ⑥ | 后台执行只读拉主图、远程原图下载、本地图片校验复制、cmhub 套图生成与AI帮写；拉图和本轮下载支持安全边界协作停止，worker 不直接操作 QWidget |
> 采集、生成、更新都是耗时操作，使用 `QObject` worker + `QThread`。Worker 不直接操作 QWidget，只通过 signal 通知主线程刷新 UI。
docs/tasks/T-642.md
+6 -2
View File
@@ -3,7 +3,7 @@ id: T-642
title: 商品套图拉取主图可停止与本轮图片清理
phase: 7
deps: [T-641]
status: TODO
status: DONE
created: 2026-07-16
---
@@ -151,4 +151,8 @@ git diff --check
## 执行记录
- 待执行。
- 2026-07-16：⑥拉取主图新增内存 `pull_run_token` 生命周期，覆盖 URL 读取和本轮远程原图下载。按钮在活动期间显示「停止拉取蝦皮」，停止请求后显示「正在停止...」；停止弹窗提供「停止并保留 / 停止并清除本次新增 / 继续拉取」。重复停止不重复弹窗，账号和商品 ID 在本轮结束前保持锁定，其他任务仍可切换使用。
- 2026-07-16：`ImageStudioPullImagesWorker` 和 `pull_remote_main_image_urls()` 增加可选 `should_stop`，在账号准备、打开商品页、读取 URL 和同步资产前后检查；当前 CDP 请求返回后安全停止，并继续复用 `close_readonly_product()` 关闭本轮只读 tab。线程结束增加 token 校验和幂等兜底；旧轮次 URL/下载迟到结果不会清除新 worker 或覆盖新状态。
- 2026-07-16：本轮下载使用 asset ID → token 归属，只取消本轮排队和活动下载，不影响用户此前手动触发的下载。远程图片 requests 下载在请求返回、分块读取和保存前检查取消并关闭 response、清理临时文件；不使用线程强杀或结束 Chrome。
- 2026-07-16：「停止并清除本次新增」按拉取前 asset 快照与本轮远程 asset 差集处理，只移除当前项目中未被 job/selection 引用的新远程原图；被引用图片保留。本地手动导入和拉取前图片不删除，旧图片状态/顺序恢复，新但被引用的图片追加到旧图片之后；本地文件和蝦皮线上图片不删除。
- 2026-07-16：拉取生命周期诊断日志只记录 token 短值、项目/商品、数量、清理策略和用时。同步更新架构、API 与路由文档；补齐安全边界取消、response 关闭、快照恢复、本轮下载隔离、按钮三态、清理/引用保护、旧 token、线程结束兜底和真实 `QThread` 完成回归测试。相关 81 项通过；商品套图 GUI 42 项在主工作区通过，其余测试文件在短路径隔离 worktree 中逐文件通过，共覆盖 550 项；Ruff、`compileall` 与 `git diff --check` 均通过。未连接真实账号执行 CDP 停止冒烟，需在测试商品 `51100639510` 上人工确认当前有界请求返回后的停止时延和自动新建 tab 关闭表现。
tests/test_image_studio.py
+89
View File
@@ -1060,6 +1060,95 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_pull_remote_main_image_urls_stops_after_open_and_closes_created_tab(self):
with self.make_temp_dir() as temp_dir:
cfg = self._config(temp_dir)
accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg)
cdp = SimpleNamespace()
stop_values = iter([False, False, False, False, True])
with mock.patch("app.image_studio.chrome.is_running", return_value=True), \
mock.patch(
"app.image_studio.accounts.detect_login",
return_value={"logged_in": True},
), \
mock.patch(
"app.image_studio.editor.open_product",
return_value=cdp,
) as open_product, \
mock.patch(
"app.image_studio.editor.read_product_image_urls",
) as read_urls, \
mock.patch(
"app.image_studio.editor.close_readonly_product",
) as close_readonly:
with self.assertRaises(image_studio.ImageStudioPullCancelled):
image_studio.pull_remote_main_image_urls(
"alias-a",
"51100639510",
path=cfg["db_path"],
config=cfg,
should_stop=lambda: next(stop_values),
)
open_product.assert_called_once()
read_urls.assert_not_called()
close_readonly.assert_called_once_with(cdp)
self.assert_removed(temp_dir)
def test_restore_original_asset_snapshot_restores_existing_state(self):
with self.make_temp_dir() 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,
)
assets = image_studio.sync_original_asset_urls(
project.id,
[
{"index": 1, "src": "https://susercontent.com/one.jpg"},
{"index": 2, "src": "https://susercontent.com/two.jpg"},
],
path=db_path,
)
snapshot = [
{
"id": assets[0].id,
"status": image_studio.ASSET_STATUS_AVAILABLE,
"source_order": 1,
},
{
"id": assets[1].id,
"status": image_studio.ASSET_STATUS_AVAILABLE,
"source_order": 2,
},
]
image_studio.sync_original_asset_urls(
project.id,
[{"index": 1, "src": "https://susercontent.com/two.jpg"}],
path=db_path,
)
restored = image_studio.restore_original_asset_snapshot(
project.id,
snapshot,
path=db_path,
)
self.assertEqual(
[
(image_studio.ASSET_STATUS_AVAILABLE, 1),
(image_studio.ASSET_STATUS_AVAILABLE, 2),
],
[(asset.status, asset.source_order) for asset in restored],
)
self.assert_removed(temp_dir)
if __name__ == "__main__":
unittest.main()
tests/test_image_studio_images.py
+31
View File
@@ -29,6 +29,7 @@ class _Response:
self.headers = dict(headers or {})
self.url = url
self.chunk_size = chunk_size
self.closed = False
def iter_content(self, chunk_size=65536):
if self.chunk_size:
@@ -37,6 +38,9 @@ class _Response:
return
yield self.content
def close(self):
self.closed = True
class ImageStudioImageTests(TempDirMixin, unittest.TestCase):
def _png_bytes(self, size=(20, 16), color=(80, 120, 200, 255)):
@@ -246,6 +250,33 @@ class ImageStudioImageTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_download_remote_image_cancels_between_chunks_and_closes_response(self):
png = self._png_bytes(size=(120, 120))
response = _Response(
content=png,
headers={"Content-Type": "image/png"},
chunk_size=32,
)
session = SimpleNamespace(get=lambda *args, **kwargs: response)
checks = {"count": 0}
def should_stop():
checks["count"] += 1
return checks["count"] >= 4
with mock.patch(
"app.image_studio_images.socket.getaddrinfo",
return_value=self._public_dns(),
):
with self.assertRaises(image_studio_images.ImageStudioImageCancelled):
image_studio_images.download_remote_image(
"https://cdn.example.com/a.png",
session=session,
should_stop=should_stop,
)
self.assertTrue(response.closed)
def test_import_original_files_copies_valid_images_deduplicates_and_limits(self):
with self.make_temp_dir() as temp_dir:
db_path = os.path.join(temp_dir, "cmshopee.db")
tests/test_product_suite_gui.py
+292
View File
@@ -1866,6 +1866,298 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_pull_button_stops_with_confirmation_and_repeated_click_is_ignored(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
tab.item_id_edit.setText("51100639510")
with mock.patch.object(tab, "_confirm", return_value=True), \
mock.patch.object(tab, "_start_thread", return_value=mock.Mock()):
tab.pull_main_images()
self.assertTrue(state.pull_running())
self.assertEqual("停止拉取蝦皮", tab.pull_button.text())
self.assertTrue(tab.pull_button.isEnabled())
worker = state.pull_worker
with mock.patch.object(tab, "_pull_stop_action", return_value="continue"):
tab.pull_main_images()
self.assertFalse(worker.is_cancelled())
self.assertFalse(state.pull_stop_requested)
with mock.patch.object(tab, "_pull_stop_action", return_value="keep"):
tab.pull_main_images()
self.assertTrue(worker.is_cancelled())
self.assertTrue(state.pull_stop_requested)
self.assertEqual("正在停止...", tab.pull_button.text())
stop_action = mock.Mock(return_value="clear_current")
with mock.patch.object(tab, "_pull_stop_action", stop_action):
tab.pull_main_images()
stop_action.assert_not_called()
state.pull_worker = None
state.pull_thread = None
tab._pull_run_states.clear()
state.pull_run_token = ""
self.assert_removed(temp_dir)
def test_stopping_pull_cancels_only_current_pull_downloads(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
pull_worker = mock.Mock()
manual_worker = mock.Mock()
state.download_queue = [11, 12]
state.downloads = {
13: (pull_worker, mock.Mock()),
14: (manual_worker, mock.Mock()),
}
state.download_tokens = {
11: "pull-token",
12: "",
13: "pull-token",
14: "",
}
state.pull_download_asset_ids = {11, 13}
tab._cancel_pull_downloads(state)
self.assertEqual([12], state.download_queue)
self.assertEqual({13}, state.pull_download_asset_ids)
self.assertNotIn(11, state.download_tokens)
pull_worker.cancel.assert_called_once()
manual_worker.cancel.assert_not_called()
self.assert_removed(temp_dir)
def test_clear_stopped_pull_removes_only_new_unreferenced_remote_assets(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
account = accounts.create_account(
"主店",
"alias-a",
debug_port=9222,
config=config,
)
project = image_studio.create_or_get_project(
account,
item_id="51100639510",
path=config["db_path"],
)
existing_remote = image_studio.sync_original_asset_urls(
project.id,
[{"index": 1, "src": "https://susercontent.com/existing.jpg"}],
path=config["db_path"],
)[0]
local_path = os.path.join(temp_dir, "local.png")
self._write_image(local_path)
local_asset = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
local_path=local_path,
source_order=2,
path=config["db_path"],
)
before_assets = image_studio.list_assets(
project.id,
kind=image_studio.ASSET_KIND_ORIGINAL,
path=config["db_path"],
)
synced = image_studio.sync_original_asset_urls(
project.id,
[
{"index": 1, "src": "https://susercontent.com/new-one.jpg"},
{"index": 2, "src": "https://susercontent.com/new-two.jpg"},
],
path=config["db_path"],
)
new_assets = [
asset
for asset in synced
if asset.remote_url
and asset.remote_url.endswith(("new-one.jpg", "new-two.jpg"))
]
referenced = next(
asset for asset in new_assets if asset.remote_url.endswith("new-two.jpg")
)
image_studio.create_job(
project.id,
source_asset_id=referenced.id,
path=config["db_path"],
)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
state.account_alias = "alias-a"
state.item_id = project.item_id
state.project_id = project.id
state.project_binding_state = project.binding_state
state.pull_run_token = "clear-pull"
state.pull_stop_requested = True
state.pull_cleanup_mode = "clear_current"
state.pull_before_asset_ids = {int(asset.id) for asset in before_assets}
state.pull_before_asset_states = [
{
"id": int(asset.id),
"status": asset.status,
"source_order": int(asset.source_order),
}
for asset in before_assets
]
state.pull_asset_ids = {int(asset.id) for asset in new_assets}
state.pull_started_at = time.monotonic()
tab._pull_run_states["clear-pull"] = state.key
messages = []
tab._message = lambda title, message, **kwargs: messages.append(
(title, message)
)
self.assertTrue(tab._finalize_pull(state, "clear-pull"))
remaining = {
asset.id: asset
for asset in image_studio.list_assets(
project.id,
kind=image_studio.ASSET_KIND_ORIGINAL,
path=config["db_path"],
)
}
removed = next(
asset for asset in new_assets if asset.id != referenced.id
)
self.assertNotIn(removed.id, remaining)
self.assertIn(referenced.id, remaining)
self.assertEqual(
image_studio.ASSET_STATUS_AVAILABLE,
remaining[existing_remote.id].status,
)
self.assertEqual(1, remaining[existing_remote.id].source_order)
self.assertIn(local_asset.id, remaining)
self.assertGreater(
remaining[referenced.id].source_order,
remaining[local_asset.id].source_order,
)
self.assertEqual("拉取蝦皮主图已停止", messages[-1][0])
self.assertIn("清理1张", messages[-1][1])
self.assertIn("因引用保留1张", messages[-1][1])
self.assert_removed(temp_dir)
def test_old_pull_result_does_not_override_current_run(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
state.pull_run_token = "current-pull"
state.pull_worker = mock.Mock()
tab._pull_run_states["current-pull"] = state.key
tab._on_pull_finished(
"old-pull",
{"project": None, "assets": [], "count": 0},
)
self.assertEqual("current-pull", state.pull_run_token)
self.assertIsNotNone(state.pull_worker)
state.pull_worker = None
state.pull_run_token = ""
tab._pull_run_states.clear()
self.assert_removed(temp_dir)
def test_pull_thread_finished_fallback_finalizes_requested_stop(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
state.pull_run_token = "pull-fallback"
state.pull_stop_requested = True
state.pull_worker = mock.Mock()
state.pull_thread = mock.Mock()
state.pull_started_at = time.monotonic()
tab._pull_run_states["pull-fallback"] = state.key
messages = []
tab._message = lambda title, message, **kwargs: messages.append(
(title, message)
)
tab._handle_pull_thread_finished("pull-fallback")
self.assertEqual("", state.pull_run_token)
self.assertIsNone(state.pull_worker)
self.assertEqual("拉取蝦皮主图已停止", messages[-1][0])
self.assertEqual("拉取蝦皮主图", tab.pull_button.text())
self.assert_removed(temp_dir)
def test_pull_real_qthread_completion_restores_button_once(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
account = accounts.create_account(
"主店",
"alias-a",
debug_port=9222,
config=config,
)
project = image_studio.create_or_get_project(
account,
item_id="51100639510",
path=config["db_path"],
)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
tab.item_id_edit.setText(project.item_id)
with mock.patch.object(tab, "_confirm", return_value=True), \
mock.patch(
"app.gui.workers.image_studio.pull_remote_main_image_urls",
return_value={
"project": project,
"assets": [],
"account": account,
},
):
tab.pull_main_images()
pull_thread = state.pull_thread
deadline = time.monotonic() + 3
while state.pull_running() and time.monotonic() < deadline:
QTest.qWait(20)
self.app.processEvents()
while pull_thread is not None and time.monotonic() < deadline:
try:
running = pull_thread.isRunning()
except RuntimeError:
pull_thread = None
break
if not running:
break
QTest.qWait(20)
self.app.processEvents()
self.assertFalse(state.pull_running())
self.assertIsNone(state.pull_worker)
self.assertEqual("拉取蝦皮主图", tab.pull_button.text())
if pull_thread is not None:
self.assertFalse(pull_thread.isRunning())
self.assert_removed(temp_dir)
def test_generation_confirmation_prevents_job_creation_and_retry_bypasses_it(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
tests/test_workers.py
+33 -2
View File
@@ -7,7 +7,7 @@ os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
from _helpers import REPO_ROOT # noqa: F401
from app import workers
from app import image_studio_images, workers
if workers.QT_IMPORT_ERROR is not None:
raise unittest.SkipTest("PySide6 未安装")
@@ -16,7 +16,10 @@ from PySide6.QtCore import QEventLoop, QTimer
from PySide6.QtWidgets import QApplication
from app.workers import BaseWorker, run_worker
from app.gui.workers import ImageStudioDownloadOriginalWorker
from app.gui.workers import (
ImageStudioDownloadOriginalWorker,
ImageStudioPullImagesWorker,
)
class DemoWorker(BaseWorker):
@@ -165,6 +168,34 @@ class WorkerTests(unittest.TestCase):
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)
if __name__ == "__main__":
unittest.main()