| from __future__ import annotations |
|
|
| import base64 |
| import json |
| import os |
| import sqlite3 |
| import tempfile |
| import time |
| import unittest |
| from contextlib import closing |
| from pathlib import Path |
| from unittest import mock |
|
|
| from services.image_task_service import ImageTaskService, _public_task |
|
|
|
|
| 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 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_internal_model_daily_limit_settles_successful_task(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| calls: dict[str, list[dict[str, object]]] = {"reserve": [], "consume": [], "refund": []} |
|
|
| class FakeAuth: |
| def reserve_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["reserve"].append({"key_id": key_id, **kwargs}) |
| return {"_internal_image_model_daily_limit_reservation": {"amount": kwargs.get("amount")}} |
|
|
| def consume_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["consume"].append({"key_id": key_id, **kwargs}) |
| return {} |
|
|
| def refund_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["refund"].append({"key_id": key_id, **kwargs}) |
| return {} |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json") |
| with mock.patch("services.auth_service.auth_service", FakeAuth()): |
| service.submit_generation( |
| OTHER_OWNER, |
| client_task_id="internal-limit-success", |
| prompt="cat", |
| model="gpt-image-2-4k", |
| size=None, |
| base_url="http://local.test", |
| internal_model_limit_user_id="owner-2", |
| internal_model_limit_model_id="gpt-image-2-4k", |
| internal_model_limit_amount=1, |
| ) |
| wait_for_task(service, OTHER_OWNER, "internal-limit-success", "success") |
|
|
| self.assertEqual(calls["reserve"], [{"key_id": "owner-2", "model_id": "gpt-image-2-4k", "amount": 1}]) |
| self.assertEqual(calls["consume"], [{"key_id": "owner-2", "model_id": "gpt-image-2-4k", "amount": 1}]) |
| self.assertEqual(calls["refund"], []) |
|
|
| def test_internal_model_daily_limit_refunds_failed_task(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| calls: dict[str, list[dict[str, object]]] = {"reserve": [], "consume": [], "refund": []} |
|
|
| class FakeAuth: |
| def reserve_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["reserve"].append({"key_id": key_id, **kwargs}) |
| return {"_internal_image_model_daily_limit_reservation": {"amount": kwargs.get("amount")}} |
|
|
| def consume_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["consume"].append({"key_id": key_id, **kwargs}) |
| return {} |
|
|
| def refund_internal_image_model_daily_limit(self, key_id: str, **kwargs): |
| calls["refund"].append({"key_id": key_id, **kwargs}) |
| return {} |
|
|
| def handler(_payload): |
| raise RuntimeError("upstream failed") |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| with mock.patch("services.auth_service.auth_service", FakeAuth()): |
| service.submit_generation( |
| OTHER_OWNER, |
| client_task_id="internal-limit-failed", |
| prompt="cat", |
| model="gpt-image-2-4k", |
| size=None, |
| base_url="http://local.test", |
| internal_model_limit_user_id="owner-2", |
| internal_model_limit_model_id="gpt-image-2-4k", |
| internal_model_limit_amount=1, |
| ) |
| wait_for_task(service, OTHER_OWNER, "internal-limit-failed", "error") |
|
|
| self.assertEqual(calls["reserve"], [{"key_id": "owner-2", "model_id": "gpt-image-2-4k", "amount": 1}]) |
| self.assertEqual(calls["consume"], []) |
| self.assertEqual(calls["refund"], [{"key_id": "owner-2", "model_id": "gpt-image-2-4k", "amount": 1}]) |
|
|
| def test_normal_pool_task_settles_snapshotted_fractional_price(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| calls: dict[str, list[dict[str, object]]] = {"reserve": [], "consume": [], "refund": []} |
|
|
| class FakeAuth: |
| def reserve_credits(self, key_id: str, amount: int | float, **kwargs): |
| calls["reserve"].append({"key_id": key_id, "amount": amount, **kwargs}) |
| return { |
| "_credit_reservation": { |
| "daily": 0, |
| "permanent_paid": amount, |
| "packages": [], |
| "grants": [], |
| } |
| } |
|
|
| def consume_reserved_credits(self, key_id: str, amount: int | float, **kwargs): |
| calls["consume"].append({"key_id": key_id, "amount": amount, **kwargs}) |
| return {} |
|
|
| def refund_reserved_credits(self, key_id: str, amount: int | float, **kwargs): |
| calls["refund"].append({"key_id": key_id, "amount": amount, **kwargs}) |
| return {} |
|
|
| pricing = { |
| "base_unit_price": 1, |
| "pay_percentage": 50, |
| "unit_price": 0.5, |
| "estimated_total": 0.5, |
| "rule_id": "weekday-offpeak", |
| "rule_name": "工作日低谷", |
| "evaluated_at": "2026-09-02T07:00:00+08:00", |
| } |
| service = self.make_service(Path(tmp_dir) / "image_tasks.json") |
| with mock.patch("services.auth_service.auth_service", FakeAuth()): |
| service.submit_generation( |
| {"id": "normal-1", "name": "Normal", "role": "normal"}, |
| client_task_id="priced-pool-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| credit_user_id="normal-1", |
| credit_amount=0.5, |
| credit_pricing=pricing, |
| ) |
| task = wait_for_task( |
| service, |
| {"id": "normal-1", "name": "Normal", "role": "normal"}, |
| "priced-pool-task", |
| "success", |
| ) |
|
|
| self.assertEqual([item["amount"] for item in calls["reserve"]], [0.5]) |
| self.assertEqual([item["amount"] for item in calls["consume"]], [0.5]) |
| self.assertEqual(calls["refund"], []) |
| with service._lock: |
| stored_task = next(item for item in service._tasks.values() if item.get("id") == "priced-pool-task") |
| self.assertEqual(stored_task["credit_pricing"]["rule_id"], "weekday-offpeak") |
|
|
| 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_success_log_records_generation_account_email(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| captured_logs: list[dict[str, object]] = [] |
|
|
| def handler(_payload): |
| return { |
| "data": [{"url": "http://example.test/image.png"}], |
| "_account_email": "plus@example.com", |
| } |
|
|
| def capture_log(_type, _summary="", detail=None, **_data): |
| captured_logs.append(dict(detail or {})) |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| with mock.patch("services.image_task_service.log_service.add", capture_log): |
| service.submit_generation( |
| OWNER, |
| client_task_id="account-log-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
| wait_for_task(service, OWNER, "account-log-task", "success") |
|
|
| self.assertTrue( |
| any(item.get("account_email") == "plus@example.com" for item in captured_logs), |
| captured_logs, |
| ) |
|
|
| def test_third_party_task_uses_url_response_format_and_logs_final_format(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| captured_payload: dict[str, object] = {} |
| captured_logs: list[dict[str, object]] = [] |
|
|
| def handler(payload): |
| captured_payload.update(payload) |
| return { |
| "data": [{"url": "http://example.test/image.png"}], |
| "_third_party_image_api": { |
| "group_name": "4k", |
| "connection_name": "API A", |
| "model_id": "gpt-image-2", |
| "upstream_model": "gpt-image-2", |
| "resolution": "4k", |
| "response_format": "url", |
| }, |
| } |
|
|
| def capture_log(_type, _summary="", detail=None, **_data): |
| captured_logs.append(dict(detail or {})) |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| with mock.patch("services.image_task_service.log_service.add", capture_log): |
| service.submit_generation( |
| OWNER, |
| client_task_id="third-party-format-task", |
| prompt="4k poster", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| resolution="4k", |
| third_party_image_group_id="group-1", |
| ) |
| wait_for_task(service, OWNER, "third-party-format-task", "success") |
|
|
| self.assertEqual(captured_payload["response_format"], "url") |
| self.assertTrue( |
| any(item.get("image_api_response_format") == "url" for item in captured_logs), |
| captured_logs, |
| ) |
|
|
| def test_inline_base64_task_result_is_saved_as_lightweight_site_url(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| encoded = base64.b64encode(b"image-bytes").decode("ascii") |
|
|
| class Stored: |
| url = "http://local.test/images/2026/08/29/result.png" |
|
|
| def handler(_payload): |
| return {"data": [{"b64_json": encoded, "revised_prompt": "cat"}]} |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| with mock.patch("services.image_task_service.image_storage_service.save", return_value=Stored()) as save: |
| service.submit_generation( |
| OWNER, |
| client_task_id="base64-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
| task = wait_for_task(service, OWNER, "base64-task", "success") |
|
|
| self.assertEqual(save.call_count, 1) |
| self.assertEqual(save.call_args.args[0], b"image-bytes") |
| self.assertEqual(save.call_args.kwargs["base_url"], "http://local.test") |
| self.assertEqual(task["data"], [{"url": Stored.url, "revised_prompt": "cat"}]) |
|
|
| def test_running_task_exposes_progress_and_elapsed_time(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| progress_ready = __import__("threading").Event() |
| finish = __import__("threading").Event() |
|
|
| def handler(payload): |
| payload["progress_callback"]("generation_success") |
| progress_ready.set() |
| finish.wait(1) |
| payload["progress_callback"]("receiving_image") |
| return {"data": [{"url": "http://example.test/image.png"}]} |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| service.submit_generation( |
| OWNER, |
| client_task_id="progress-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
|
|
| self.assertTrue(progress_ready.wait(1)) |
| running = service.list_tasks(OWNER, ["progress-task"])["items"][0] |
| self.assertEqual(running["status"], "running") |
| self.assertEqual(running["progress"], "generation_success") |
| self.assertEqual([item["step"] for item in running["progress_events"]], ["generation_success"]) |
| self.assertIsInstance(running["elapsed_secs"], float) |
|
|
| finish.set() |
| completed = wait_for_task(service, OWNER, "progress-task", "success") |
| self.assertEqual(completed["duration_ms"] >= 0, True) |
| self.assertEqual( |
| [item["step"] for item in completed["progress_events"]], |
| ["generation_success", "receiving_image"], |
| ) |
|
|
| def test_running_task_elapsed_time_never_goes_negative(self): |
| public_task = _public_task({ |
| "id": "future-task", |
| "status": "running", |
| "mode": "generate", |
| "model": "gpt-image-2", |
| "size": "", |
| "created_at": "2026-08-03 12:00:00", |
| "started_at": "2026-08-03 12:00:00", |
| "updated_at": "2026-08-03 12:00:00", |
| "started_ts": time.time() + 8 * 60 * 60, |
| }) |
|
|
| self.assertEqual(public_task["elapsed_secs"], 0.0) |
|
|
| def test_public_task_redacts_account_email_from_error_without_mutating_admin_data(self): |
| email = "private-task@example.com" |
| task = { |
| "id": "private-error-task", |
| "status": "error", |
| "mode": "generate", |
| "model": "gpt-image-2", |
| "size": "", |
| "created_at": "2026-08-03 12:00:00", |
| "started_at": "2026-08-03 12:00:00", |
| "updated_at": "2026-08-03 12:00:00", |
| "account_email": email, |
| "error": f'HTTP 400: {{"error":{{"account_email":"{email}"}}}}', |
| } |
|
|
| public_task = _public_task(task) |
|
|
| self.assertNotIn(email, public_task["error"]) |
| self.assertNotIn("account_email", public_task) |
| self.assertEqual(task["account_email"], email) |
| self.assertIn(email, task["error"]) |
|
|
| def test_memory_reset_is_passed_to_edit_handler(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| captured_payload = {} |
|
|
| def handler(payload): |
| captured_payload.update(payload) |
| return {"data": [{"url": "http://example.test/edit.png"}]} |
|
|
| service = self.make_service(Path(tmp_dir) / "image_tasks.json", handler) |
| service.submit_edit( |
| OWNER, |
| client_task_id="memory-reset-task", |
| prompt="increase saturation", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| images=[(b"image", "reference.png", "image/png")], |
| memory_enabled=True, |
| memory_conversation_id="local-conversation", |
| memory_reset=True, |
| ) |
| wait_for_task(service, OWNER, "memory-reset-task", "success") |
|
|
| self.assertTrue(captured_payload["memory"]) |
| self.assertTrue(captured_payload["memory_reset"]) |
| self.assertEqual(captured_payload["memory_conversation_id"], "local-conversation") |
|
|
| 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.assertTrue(path.with_suffix(".db").is_file()) |
| self.assertEqual([item["status"] for item in result["items"]], ["error", "error"]) |
| self.assertTrue(all("已中断" in item.get("error", "") for item in result["items"])) |
|
|
| def test_startup_quarantines_corrupt_database_and_rebuilds_from_json(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| path = Path(tmp_dir) / "image_tasks.json" |
| path.write_text( |
| json.dumps( |
| { |
| "tasks": [ |
| { |
| "id": "legacy-task", |
| "owner_id": "owner-1", |
| "status": "success", |
| "mode": "generate", |
| "model": "gpt-image-2", |
| "created_at": "2026-08-25 10:00:00", |
| "updated_at": "2026-08-25 10:00:01", |
| "data": [{"url": "https://example.test/recovered.png"}], |
| } |
| ] |
| } |
| ), |
| encoding="utf-8", |
| ) |
| database_path = path.with_suffix(".db") |
| database_path.write_bytes(b"not a sqlite database") |
|
|
| service = self.make_service(path) |
| result = service.list_tasks(OWNER, ["legacy-task"]) |
|
|
| self.assertEqual(result["items"][0]["status"], "success") |
| self.assertTrue(list(Path(tmp_dir).glob("image_tasks.db.corrupt-*"))) |
| with closing(sqlite3.connect(database_path)) as connection: |
| self.assertEqual(connection.execute("PRAGMA integrity_check").fetchone()[0], "ok") |
| self.assertEqual(connection.execute("PRAGMA journal_mode").fetchone()[0], "delete") |
|
|
| def test_environment_error_relocates_database(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| |
| blocker = Path(tmp_dir) / "blocker" |
| blocker.write_text("not a directory", encoding="utf-8") |
| fallback = Path(tmp_dir) / "fallback" / "image_tasks.db" |
| path = blocker / "image_tasks.json" |
|
|
| with mock.patch.dict(os.environ, {"CHATGPT2API_IMAGE_TASKS_DB": str(fallback)}): |
| service = self.make_service(path) |
| self.assertEqual(service.db_path, fallback) |
| service.submit_generation( |
| OWNER, |
| client_task_id="relocated-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
| task = wait_for_task(service, OWNER, "relocated-task", "success") |
| self.assertEqual(task["data"][0]["url"], "http://example.test/image.png") |
|
|
| self.assertTrue(fallback.is_file()) |
| with closing(sqlite3.connect(fallback)) as connection: |
| self.assertEqual(connection.execute("PRAGMA integrity_check").fetchone()[0], "ok") |
|
|
| def test_environment_error_falls_back_to_memory(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| blocker = Path(tmp_dir) / "blocker" |
| blocker.write_text("not a directory", encoding="utf-8") |
| path = blocker / "image_tasks.json" |
|
|
| with mock.patch.object( |
| ImageTaskService, "_database_fallback_candidates", return_value=[] |
| ): |
| service = self.make_service(path) |
| original_db_path = service.db_path |
| self.assertIsNotNone(service._memory_database) |
| self.assertEqual(service.db_path, original_db_path) |
|
|
| service.submit_generation( |
| OWNER, |
| client_task_id="memory-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
| task = wait_for_task(service, OWNER, "memory-task", "success") |
| self.assertEqual(task["data"][0]["url"], "http://example.test/image.png") |
| self.assertFalse(original_db_path.exists()) |
|
|
| def test_runtime_readonly_error_relocates_and_retries_write(self): |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| path = Path(tmp_dir) / "image_tasks.json" |
| fallback = Path(tmp_dir) / "fallback" / "image_tasks.db" |
| service = self.make_service(path) |
| original_writer = service._write_database_task_list_locked |
| failed_once = False |
|
|
| def fail_primary_once(task_list, *, replace_all=False): |
| nonlocal failed_once |
| if not failed_once: |
| failed_once = True |
| raise sqlite3.OperationalError("attempt to write a readonly database") |
| return original_writer(task_list, replace_all=replace_all) |
|
|
| with ( |
| mock.patch.dict(os.environ, {"CHATGPT2API_IMAGE_TASKS_DB": str(fallback)}), |
| mock.patch.object(service, "_write_database_task_list_locked", side_effect=fail_primary_once), |
| ): |
| service.submit_generation( |
| OWNER, |
| client_task_id="runtime-relocated-task", |
| prompt="cat", |
| model="gpt-image-2", |
| size=None, |
| base_url="http://local.test", |
| ) |
| task = wait_for_task(service, OWNER, "runtime-relocated-task", "success") |
|
|
| self.assertEqual(service.db_path, fallback) |
| self.assertEqual(task["data"][0]["url"], "http://example.test/image.png") |
| with closing(sqlite3.connect(fallback)) as connection: |
| stored = connection.execute( |
| "SELECT payload FROM image_tasks WHERE task_id = ?", |
| ("runtime-relocated-task",), |
| ).fetchone() |
| self.assertIsNotNone(stored) |
| self.assertEqual(json.loads(stored[0])["status"], "success") |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
