182 lines
7.4 KiB
Python
182 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
import time
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from services.image_task_service import ImageTaskService
|
|
|
|
|
|
OWNER = {"id": "owner-1", "name": "Owner", "role": "admin"}
|
|
OTHER_OWNER = {"id": "owner-2", "name": "Other", "role": "user"}
|
|
|
|
|
|
def wait_for_task(service: ImageTaskService, identity: dict[str, object], task_id: str, status: str, timeout: float = 2.0):
|
|
deadline = time.time() + timeout
|
|
last = None
|
|
while time.time() < deadline:
|
|
result = service.list_tasks(identity, [task_id])
|
|
last = (result.get("items") or [None])[0]
|
|
if last and last.get("status") == status:
|
|
return last
|
|
time.sleep(0.02)
|
|
raise AssertionError(f"task {task_id} did not reach {status}, last={last}")
|
|
|
|
|
|
class ImageTaskServiceTests(unittest.TestCase):
|
|
def test_elapsed_time_includes_startup_and_stage_timings_survive_reload(self):
|
|
entered = threading.Event()
|
|
proceed = threading.Event()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = Path(tmp_dir) / "tasks.json"
|
|
|
|
def handler(payload):
|
|
payload["progress_callback"]("starting_generation")
|
|
entered.set()
|
|
if not proceed.wait(timeout=2):
|
|
raise RuntimeError("test did not release task")
|
|
payload["progress_callback"]("image_stream_resolve_start")
|
|
return {"data": [{"url": "http://example.test/image.png"}]}
|
|
|
|
service = self.make_service(path, handler)
|
|
service.submit_generation(OWNER, client_task_id="timed", prompt="cat", model="gpt-image-2", size=None, base_url="http://local.test")
|
|
try:
|
|
self.assertTrue(entered.wait(timeout=2))
|
|
running = service.list_tasks(OWNER, ["timed"])["items"][0]
|
|
self.assertIn("elapsed_secs", running)
|
|
self.assertEqual(running["progress"], "starting_generation")
|
|
started_ts = next(iter(service._tasks.values()))["started_ts"]
|
|
finally:
|
|
proceed.set()
|
|
task = wait_for_task(service, OWNER, "timed", "success")
|
|
self.assertEqual(next(iter(service._tasks.values()))["started_ts"], started_ts)
|
|
self.assertIn("starting_generation", task["stage_durations_ms"])
|
|
self.assertIn("image_stream_resolve_start", task["stage_durations_ms"])
|
|
reloaded = self.make_service(path).list_tasks(OWNER, ["timed"])["items"][0]
|
|
self.assertEqual(reloaded["stage_durations_ms"], task["stage_durations_ms"])
|
|
|
|
def make_service(self, path: Path, handler=None) -> ImageTaskService:
|
|
return ImageTaskService(
|
|
path,
|
|
generation_handler=handler or (lambda _payload: {"data": [{"url": "http://example.test/image.png"}]}),
|
|
edit_handler=handler or (lambda _payload: {"data": [{"url": "http://example.test/edit.png"}]}),
|
|
retention_days_getter=lambda: 30,
|
|
)
|
|
|
|
def test_duplicate_submit_uses_existing_task(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
calls = 0
|
|
|
|
def handler(_payload):
|
|
nonlocal calls
|
|
calls += 1
|
|
time.sleep(0.05)
|
|
return {"data": [{"url": "http://example.test/image.png"}]}
|
|
|
|
service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler)
|
|
first = service.submit_generation(
|
|
OWNER,
|
|
client_task_id="task-1",
|
|
prompt="cat",
|
|
model="gpt-image-2",
|
|
size=None,
|
|
base_url="http://local.test",
|
|
)
|
|
second = service.submit_generation(
|
|
OWNER,
|
|
client_task_id="task-1",
|
|
prompt="cat",
|
|
model="gpt-image-2",
|
|
size=None,
|
|
base_url="http://local.test",
|
|
)
|
|
|
|
self.assertEqual(first["id"], "task-1")
|
|
self.assertEqual(second["id"], "task-1")
|
|
task = wait_for_task(service, OWNER, "task-1", "success")
|
|
self.assertEqual(task["data"][0]["url"], "http://example.test/image.png")
|
|
self.assertEqual(calls, 1)
|
|
|
|
def test_different_owner_cannot_query_task(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
service = self.make_service(Path(tmp_dir) / "image_tasks.json")
|
|
service.submit_generation(
|
|
OWNER,
|
|
client_task_id="private-task",
|
|
prompt="cat",
|
|
model="gpt-image-2",
|
|
size=None,
|
|
base_url="http://local.test",
|
|
)
|
|
|
|
wait_for_task(service, OWNER, "private-task", "success")
|
|
result = service.list_tasks(OTHER_OWNER, ["private-task"])
|
|
|
|
self.assertEqual(result["items"], [])
|
|
self.assertEqual(result["missing_ids"], ["private-task"])
|
|
|
|
def test_success_task_persists_to_new_service_instance(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = Path(tmp_dir) / "image_tasks.json"
|
|
service = self.make_service(path)
|
|
service.submit_generation(
|
|
OWNER,
|
|
client_task_id="persisted-task",
|
|
prompt="cat",
|
|
model="gpt-image-2",
|
|
size=None,
|
|
base_url="http://local.test",
|
|
)
|
|
wait_for_task(service, OWNER, "persisted-task", "success")
|
|
|
|
reloaded = self.make_service(path)
|
|
result = reloaded.list_tasks(OWNER, ["persisted-task"])
|
|
|
|
self.assertEqual(result["missing_ids"], [])
|
|
self.assertEqual(result["items"][0]["status"], "success")
|
|
self.assertEqual(result["items"][0]["data"][0]["url"], "http://example.test/image.png")
|
|
|
|
def test_startup_marks_unfinished_tasks_as_error(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = Path(tmp_dir) / "image_tasks.json"
|
|
path.write_text(
|
|
json.dumps(
|
|
{
|
|
"tasks": [
|
|
{
|
|
"id": "queued-task",
|
|
"owner_id": "owner-1",
|
|
"status": "queued",
|
|
"mode": "generate",
|
|
"model": "gpt-image-2",
|
|
"created_at": "2099-01-01 00:00:00",
|
|
"updated_at": "2099-01-01 00:00:00",
|
|
},
|
|
{
|
|
"id": "running-task",
|
|
"owner_id": "owner-1",
|
|
"status": "running",
|
|
"mode": "generate",
|
|
"model": "gpt-image-2",
|
|
"created_at": "2099-01-01 00:00:00",
|
|
"updated_at": "2099-01-01 00:00:00",
|
|
},
|
|
]
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
service = self.make_service(path)
|
|
result = service.list_tasks(OWNER, ["queued-task", "running-task"])
|
|
|
|
self.assertEqual([item["status"] for item in result["items"]], ["error", "error"])
|
|
self.assertTrue(all("已中断" in item.get("error", "") for item in result["items"]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|