T-564 async cmhub image tasks

tests/test_ai.py

@@ -1005,6 +1005,19 @@ class AITests(TempDirMixin, unittest.TestCase):
             downloads = []
 
             def fake_request(method, url, **kwargs):
+                if str(method).upper() == "GET":
+                    task_id = url.rsplit("/", 1)[-1]
+                    index = int(task_id.rsplit("-", 1)[-1])
+                    return _RequestsResponse(
+                        {
+                            "task_id": task_id,
+                            "status": "succeeded",
+                            "result": {
+                                "image_url": "https://cdn.example.com/generated-%s.png"
+                                % index
+                            },
+                        }
+                    )
                 with lock:
                     counters["request_count"] += 1
                     request_index = counters["request_count"]
@@ -1018,8 +1031,8 @@ class AITests(TempDirMixin, unittest.TestCase):
                 try:
                     return _RequestsResponse(
                         {
-                            "image_url": "https://cdn.example.com/generated-%s.png"
-                            % request_index
+                            "task_id": "cmhub-task-%s" % request_index,
+                            "status": "queued",
                         }
                     )
                 finally:
@@ -1082,10 +1095,16 @@ class AITests(TempDirMixin, unittest.TestCase):
             downloads_seen = []
 
             def fake_request(method, url, **kwargs):
+                if str(method).upper() == "GET":
+                    return _RequestsResponse(
+                        {
+                            "task_id": "cmhub-task-1",
+                            "status": "succeeded",
+                            "result": {"image_url": "https://cdn.example.com/generated.png"},
+                        }
+                    )
                 requests_seen.append((method, url, kwargs))
-                return _RequestsResponse(
-                    {"image_url": "https://cdn.example.com/generated.png"}
-                )
+                return _RequestsResponse({"task_id": "cmhub-task-1", "status": "queued"})
 
             def fake_get(url, **kwargs):
                 downloads_seen.append((url, kwargs))
@@ -1115,6 +1134,235 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_generate_batch_cmhub_async_submit_persists_task_id_before_poll(self):
+        try:
+            from PIL import Image  # noqa: F401
+        except ImportError:
+            self.skipTest("Pillow not installed")
+
+        with self.make_temp_dir() as temp_dir:
+            cfg, key_path = self._cmhub_config(temp_dir)
+            cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
+            cfg["image_dir"] = os.path.join(temp_dir, "images")
+            cfg["ai"]["generate_cover"] = True
+            batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
+            self._write_old_cover_files(tasks)
+            db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
+            task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
+            generated_png = self._png_bytes()
+            public_dns = [
+                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
+            ]
+            poll_db_values = []
+            headers_seen = []
+
+            def fake_request(method, url, **kwargs):
+                if str(method).upper() == "POST":
+                    headers_seen.append(dict(kwargs["headers"]))
+                    return _RequestsResponse(
+                        {
+                            "task_id": "cmhub-task-1",
+                            "status": "queued",
+                            "points_cost": 2,
+                            "points_balance": 80,
+                            "call_id": "call-image-1",
+                        },
+                        status_code=202,
+                    )
+                poll_db_values.append(db.get_task(task.id, path=cfg["db_path"]).image_task_id)
+                return _RequestsResponse(
+                    {
+                        "task_id": "cmhub-task-1",
+                        "status": "succeeded",
+                        "result": {"image_url": "https://cdn.example.com/generated.png"},
+                    }
+                )
+
+            def fake_get(url, **kwargs):
+                return _RequestsResponse(content=generated_png)
+
+            with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
+                    mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
+                    mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
+                summary = ai.generate_batch(
+                    [task],
+                    {"title": "标题提示", "cover": "封面 {新标题}"},
+                    ai_cfg={
+                        "config": cfg,
+                        "db_path": cfg["db_path"],
+                        "cmhub_config_path": key_path,
+                    },
+                )
+
+            self.assertTrue(summary["ok"])
+            self.assertEqual(["cmhub-task-1"], poll_db_values)
+            self.assertTrue(headers_seen[0]["Idempotency-Key"].startswith("cmshopee-task-"))
+            self.assertTrue(headers_seen[0]["X-Client-Version"])
+            updated = db.get_task(task.id, path=cfg["db_path"])
+            self.assertEqual("cmhub-task-1", updated.image_task_id)
+            self.assertTrue(updated.image_task_key)
+            self.assertTrue(os.path.exists(updated.new_cover_path))
+
+        self.assert_removed(temp_dir)
+
+    def test_generate_batch_cmhub_resumes_existing_image_task_without_submit(self):
+        try:
+            from PIL import Image  # noqa: F401
+        except ImportError:
+            self.skipTest("Pillow not installed")
+
+        with self.make_temp_dir() as temp_dir:
+            cfg, key_path = self._cmhub_config(temp_dir)
+            cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
+            cfg["image_dir"] = os.path.join(temp_dir, "images")
+            cfg["ai"]["generate_cover"] = True
+            batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
+            self._write_old_cover_files(tasks)
+            db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
+            key = db.ensure_image_task_key(tasks[0].id, path=cfg["db_path"])
+            db.set_image_task_submitted(tasks[0].id, "cmhub-task-resume", key, path=cfg["db_path"])
+            task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
+            generated_png = self._png_bytes()
+            public_dns = [
+                (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
+            ]
+            posts = []
+
+            def fake_request(method, url, **kwargs):
+                if str(method).upper() == "POST":
+                    posts.append((method, url, kwargs))
+                    return _RequestsResponse({"task_id": "unexpected", "status": "queued"})
+                return _RequestsResponse(
+                    {
+                        "task_id": "cmhub-task-resume",
+                        "status": "succeeded",
+                        "result": {"image_url": "https://cdn.example.com/resume.png"},
+                    }
+                )
+
+            def fake_get(url, **kwargs):
+                return _RequestsResponse(content=generated_png)
+
+            with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
+                    mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
+                    mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
+                summary = ai.generate_batch(
+                    [task],
+                    {"title": "标题提示", "cover": "封面 {新标题}"},
+                    ai_cfg={
+                        "config": cfg,
+                        "db_path": cfg["db_path"],
+                        "cmhub_config_path": key_path,
+                    },
+                )
+
+            self.assertTrue(summary["ok"])
+            self.assertEqual([], posts)
+            updated = db.get_task(task.id, path=cfg["db_path"])
+            self.assertEqual("cmhub-task-resume", updated.image_task_id)
+            self.assertEqual(key, updated.image_task_key)
+            self.assertTrue(os.path.exists(updated.new_cover_path))
+
+        self.assert_removed(temp_dir)
+
+    def test_generate_batch_cmhub_failed_task_clears_image_task_state(self):
+        try:
+            from PIL import Image  # noqa: F401
+        except ImportError:
+            self.skipTest("Pillow not installed")
+
+        with self.make_temp_dir() as temp_dir:
+            cfg, key_path = self._cmhub_config(temp_dir)
+            cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
+            cfg["image_dir"] = os.path.join(temp_dir, "images")
+            cfg["ai"]["generate_cover"] = True
+            batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
+            self._write_old_cover_files(tasks)
+            db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
+            key = db.ensure_image_task_key(tasks[0].id, path=cfg["db_path"])
+            db.set_image_task_submitted(tasks[0].id, "cmhub-task-failed", key, path=cfg["db_path"])
+            task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
+
+            def fake_request(method, url, **kwargs):
+                return _RequestsResponse(
+                    {
+                        "task_id": "cmhub-task-failed",
+                        "status": "failed",
+                        "error": {"code": "upstream_timeout", "message": "上游超时"},
+                    }
+                )
+
+            with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
+                summary = ai.generate_batch(
+                    [task],
+                    {"title": "标题提示", "cover": "封面 {新标题}"},
+                    ai_cfg={
+                        "config": cfg,
+                        "db_path": cfg["db_path"],
+                        "cmhub_config_path": key_path,
+                    },
+                )
+
+            self.assertFalse(summary["ok"])
+            self.assertEqual(1, summary["failed"])
+            updated = db.get_task(task.id, path=cfg["db_path"])
+            self.assertIsNone(updated.image_task_id)
+            self.assertIsNone(updated.image_task_key)
+            self.assertIn("上游生成超时", updated.last_error)
+
+        self.assert_removed(temp_dir)
+
+    def test_generate_batch_cmhub_cancelled_poll_keeps_image_task_state(self):
+        try:
+            from PIL import Image  # noqa: F401
+        except ImportError:
+            self.skipTest("Pillow not installed")
+
+        with self.make_temp_dir() as temp_dir:
+            cfg, key_path = self._cmhub_config(temp_dir)
+            cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
+            cfg["image_dir"] = os.path.join(temp_dir, "images")
+            cfg["ai"]["generate_cover"] = True
+            batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
+            self._write_old_cover_files(tasks)
+            db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
+            task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
+            events = []
+
+            def fake_request(method, url, **kwargs):
+                if str(method).upper() == "POST":
+                    return _RequestsResponse({"task_id": "cmhub-task-cancel", "status": "queued"})
+                return _RequestsResponse({"task_id": "cmhub-task-cancel", "status": "running"})
+
+            def should_stop():
+                current = db.get_task(task.id, path=cfg["db_path"])
+                return bool(current and current.image_task_id)
+
+            with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request):
+                summary = ai.generate_batch(
+                    [task],
+                    {"title": "标题提示", "cover": "封面 {新标题}"},
+                    ai_cfg={
+                        "config": cfg,
+                        "db_path": cfg["db_path"],
+                        "cmhub_config_path": key_path,
+                        "on_event": events.append,
+                    },
+                    should_stop=should_stop,
+                )
+
+            self.assertFalse(summary["ok"])
+            self.assertTrue(summary["cancelled"])
+            updated = db.get_task(task.id, path=cfg["db_path"])
+            self.assertEqual("cmhub-task-cancel", updated.image_task_id)
+            self.assertTrue(updated.image_task_key)
+            self.assertIsNone(updated.new_cover_path)
+            self.assertTrue(
+                any("服务端任务可能仍在完成" in event.get("detail", "") for event in events)
+            )
+
+        self.assert_removed(temp_dir)
+
     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)

tests/test_db.py

@@ -29,6 +29,11 @@ class DbTests(TempDirMixin, unittest.TestCase):
                     ).fetchall()
                 }
                 self.assertTrue({"batches", "accounts", "tasks"}.issubset(tables))
+                task_columns = {
+                    row["name"]
+                    for row in conn.execute("PRAGMA table_info(tasks)").fetchall()
+                }
+                self.assertTrue({"image_task_id", "image_task_key"}.issubset(task_columns))
             finally:
                 conn.close()
 
@@ -79,6 +84,8 @@ class DbTests(TempDirMixin, unittest.TestCase):
             task = db.list_tasks(batch_id=batch_id, alias="alias", path=db_path)[0]
             self.assertEqual("imported", task.stage)
             self.assertEqual("pending", task.status)
+            self.assertIsNone(task.image_task_id)
+            self.assertIsNone(task.image_task_key)
 
             db.mark_running(task.id, "collect", path=db_path)
             self.assertEqual("running", db.list_tasks(path=db_path)[0].status)
@@ -126,6 +133,65 @@ class DbTests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_image_task_helpers_and_reset_lifecycle(self):
+        with self.make_temp_dir() 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": 2,
+                        "account_name": "shop",
+                        "alias": "alias",
+                        "item_id": "51100639510",
+                    }
+                ],
+                path=db_path,
+            )
+            task = db.list_tasks(batch_id=batch_id, path=db_path)[0]
+            db.set_collected(task.id, "旧标题", "old.jpg", path=db_path)
+            db.set_generated(task.id, "新标题", "new.jpg", path=db_path)
+
+            key = db.ensure_image_task_key(task.id, path=db_path)
+            self.assertTrue(key.startswith(f"cmshopee-task-{task.id}-"))
+            self.assertEqual(key, db.ensure_image_task_key(task.id, path=db_path))
+
+            db.set_image_task_submitted(task.id, "cmhub-task-1", key, path=db_path)
+            submitted = db.get_task(task.id, path=db_path)
+            self.assertEqual("cmhub-task-1", submitted.image_task_id)
+            self.assertEqual(key, submitted.image_task_key)
+
+            title_only = db.reset_generated(
+                task.id,
+                reset_title=True,
+                reset_cover=False,
+                path=db_path,
+            )["after"]
+            self.assertEqual("cmhub-task-1", title_only.image_task_id)
+            self.assertEqual(key, title_only.image_task_key)
+
+            cover_reset = db.reset_generated(
+                task.id,
+                reset_title=False,
+                reset_cover=True,
+                path=db_path,
+            )["after"]
+            self.assertIsNone(cover_reset.image_task_id)
+            self.assertIsNone(cover_reset.image_task_key)
+
+            new_key = db.ensure_image_task_key(task.id, path=db_path)
+            self.assertNotEqual(key, new_key)
+            db.clear_image_task(task.id, path=db_path)
+            cleared = db.get_task(task.id, path=db_path)
+            self.assertIsNone(cleared.image_task_id)
+            self.assertIsNone(cleared.image_task_key)
+
+        self.assert_removed(temp_dir)
+
     def test_mark_failed_prepends_chinese_step_without_duplicate_prefix(self):
         with self.make_temp_dir() as temp_dir:
             db_path = os.path.join(temp_dir, "cmshopee.db")