feat: AI生图功能丰富

server/tests/test_job_worker.py

@@ -7,6 +7,8 @@ from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch
 
+import httpx
+
 SERVER_DIR = Path(__file__).resolve().parents[1]
 if str(SERVER_DIR) not in sys.path:
     sys.path.insert(0, str(SERVER_DIR))
@@ -58,6 +60,46 @@ class WorkerCancellationTests(unittest.IsolatedAsyncioTestCase):
         cancelled.assert_awaited_once_with(claim.item_id)
         waiting.assert_not_awaited()
 
+    async def test_provider_poll_timeout_keeps_existing_task_for_retry(self) -> None:
+        claim = ClaimedItem(
+            job_id="00000000-0000-0000-0000-000000000001",
+            item_id="00000000-0000-0000-0000-000000000002",
+            job_kind="image_edit",
+            product_id=None,
+            provider="pcore",
+            model="gpt-image-2",
+            platform="ozon",
+            lang="ru",
+            ratio="3:4",
+            style_set=5,
+            request_payload={"context": {"title": "test"}, "ref_images": []},
+            item_payload={"kind": "custom", "title": "AI生图"},
+            provider_task_id="provider-task",
+            provider_task_kind="edits",
+        )
+        settings = SimpleNamespace(
+            image_provider="pcore",
+            dashscope_model="wan2.7-image-pro",
+            rightapi_image_model="gpt-image-2",
+            ark_image_model="doubao-seedream-4-5-251128",
+        )
+
+        with (
+            patch("services.job_worker.get_settings", return_value=settings),
+            patch("services.job_worker.build_context", return_value={}),
+            patch("services.job_worker.build_prompt", return_value="prompt"),
+            patch("services.job_worker._refs_for_job", return_value=[]),
+            patch("services.job_worker.advance_generation", new=AsyncMock(side_effect=httpx.ConnectTimeout("timeout"))),
+            patch("services.job_worker.mark_waiting_provider", new=AsyncMock()) as waiting,
+            patch("services.job_worker.mark_item_failed", new=AsyncMock()) as failed,
+        ):
+            await process_item(claim)
+
+        waiting.assert_awaited_once_with(
+            claim.item_id, "provider-task", "edits", delay_seconds=30,
+        )
+        failed.assert_not_awaited()
+
 
 if __name__ == "__main__":
     unittest.main()