64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
"""后台生图 Worker 的取消竞态回归测试。"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
SERVER_DIR = Path(__file__).resolve().parents[1]
|
|
if str(SERVER_DIR) not in sys.path:
|
|
sys.path.insert(0, str(SERVER_DIR))
|
|
|
|
from services.generator import GenerationAdvance
|
|
from services.job_worker import process_item
|
|
from services.jobs import ClaimedItem
|
|
|
|
|
|
class WorkerCancellationTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_cancelled_task_is_not_returned_to_provider_waiting_queue(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(return_value=GenerationAdvance(
|
|
provider_task_id="provider-task", provider_task_kind="edits",
|
|
))),
|
|
patch("services.job_worker.job_can_write", new=AsyncMock(return_value=False)),
|
|
patch("services.job_worker.mark_item_cancelled", new=AsyncMock()) as cancelled,
|
|
patch("services.job_worker.mark_waiting_provider", new=AsyncMock()) as waiting,
|
|
):
|
|
await process_item(claim)
|
|
|
|
cancelled.assert_awaited_once_with(claim.item_id)
|
|
waiting.assert_not_awaited()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|