from __future__ import annotations import base64 import unittest from unittest import mock from fastapi import FastAPI from fastapi.testclient import TestClient import api.image_tasks as image_tasks_module from services.third_party_image_service import ThirdPartyRoute AUTH_HEADERS = {"Authorization": "Bearer chatgpt2api"} PNG_BYTES = b"\x89PNG\r\n\x1a\n" DATA_IMAGE_URL = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode('ascii')}" class FakeImageTaskService: def __init__(self): self.generation_calls = [] self.edit_calls = [] def submit_generation(self, identity, **kwargs): self.generation_calls.append((identity, kwargs)) return { "id": kwargs["client_task_id"], "status": "success", "mode": "generate", "created_at": "2026-01-01 00:00:00", "updated_at": "2026-01-01 00:00:00", "data": [{"url": f"{kwargs['base_url']}/images/fake.png"}], } def submit_edit(self, identity, **kwargs): self.edit_calls.append((identity, kwargs)) return { "id": kwargs["client_task_id"], "status": "queued", "mode": "edit", "created_at": "2026-01-01 00:00:00", "updated_at": "2026-01-01 00:00:00", } def list_tasks(self, _identity, ids): return { "items": [ { "id": task_id, "status": "success", "mode": "generate", "created_at": "2026-01-01 00:00:00", "updated_at": "2026-01-01 00:00:00", "data": [{"url": "http://testserver/images/fake.png"}], } for task_id in ids if task_id != "missing" ], "missing_ids": [task_id for task_id in ids if task_id == "missing"], } class FakeThirdPartyImageService: def __init__(self): self.find_route_calls: list[dict[str, object]] = [] def find_route(self, model, **kwargs): self.find_route_calls.append(dict(kwargs)) return ThirdPartyRoute( group={"id": "g1", "billing_mode": "general_multiplier", "default_cost": 1}, model={"id": str(model), "upstream_model": str(model), "resolution": kwargs.get("resolution") or "1k", "default_cost": 1}, ) def assert_access(self, _identity, _route): return None def cost_for_route(self, _route, _quality, _image_count): return 1 class FakeFreeNamedThirdPartyImageService: def find_named_route(self, model, **_kwargs): return ThirdPartyRoute( group={"id": "g-free", "billing_mode": "general_multiplier", "default_cost": 1, "multiplier_enabled": True}, model={"id": str(model), "upstream_model": "gpt-image-2", "resolution": "1k", "default_cost": 0}, ) def find_route(self, model, **kwargs): return self.find_named_route(model, **kwargs) def assert_access(self, _identity, _route): return None def cost_for_route(self, _route, _quality, _image_count): return 0 class ImageTasksApiTests(unittest.TestCase): def setUp(self): self.fake_service = FakeImageTaskService() self.identity = {"id": "admin", "role": "admin", "name": "Admin", "account_pool_enabled": True} self.service_patcher = mock.patch.object(image_tasks_module, "image_task_service", self.fake_service) self.identity_patcher = mock.patch.object(image_tasks_module, "require_identity", lambda _authorization: self.identity) self.service_patcher.start() self.identity_patcher.start() self.addCleanup(self.service_patcher.stop) self.addCleanup(self.identity_patcher.stop) app = FastAPI() app.include_router(image_tasks_module.create_router()) self.client = TestClient(app) def test_create_generation_task(self): response = self.client.post( "/api/image-tasks/generations", headers=AUTH_HEADERS, json={"client_task_id": "task-1", "prompt": "cat", "model": "gpt-image-2"}, ) self.assertEqual(response.status_code, 200, response.text) payload = response.json() self.assertEqual(payload["id"], "task-1") self.assertEqual(payload["status"], "success") self.assertEqual(len(self.fake_service.generation_calls), 1) self.assertEqual(self.fake_service.generation_calls[0][1]["account_pool_strategy"], "own_first") response = self.client.post( "/api/image-tasks/generations", headers=AUTH_HEADERS, json={ "client_task_id": "task-2", "prompt": "cat", "model": "gpt-image-2", "account_pool_strategy": "shared_first", }, ) self.assertEqual(response.status_code, 200, response.text) self.assertEqual(self.fake_service.generation_calls[-1][1]["account_pool_strategy"], "shared_first") def test_disabled_personal_pool_forces_admin_pool(self): disabled_identity = { "id": "normal-1", "role": "normal", "name": "Normal", "account_pool_enabled": False, } self.identity = disabled_identity response = self.client.post( "/api/image-tasks/generations", headers=AUTH_HEADERS, json={ "client_task_id": "task-disabled", "prompt": "cat", "model": "gpt-image-2", "account_pool_strategy": "shared_first", }, ) self.assertEqual(response.status_code, 200, response.text) self.assertEqual(self.fake_service.generation_calls[-1][1]["account_pool_strategy"], "admin_only") def test_normal_pool_task_receives_submit_time_price_snapshot(self): self.identity = { "id": "normal-1", "role": "normal", "name": "Normal", "account_pool_enabled": True, } quote = { "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", } with mock.patch.object(image_tasks_module.pool_pricing_service, "quote", return_value=quote): response = self.client.post( "/api/image-tasks/generations", headers=AUTH_HEADERS, json={"client_task_id": "priced-task", "prompt": "cat", "model": "gpt-image-2"}, ) self.assertEqual(response.status_code, 200, response.text) payload = self.fake_service.generation_calls[-1][1] self.assertEqual(payload["credit_user_id"], "normal-1") self.assertEqual(payload["credit_amount"], 0.5) self.assertEqual(payload["credit_pricing"], quote) def test_create_edit_task_accepts_multiple_images(self): """测试图片编辑任务接口支持多个上传图片。""" response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={"client_task_id": "edit-1", "prompt": "edit", "model": "gpt-image-2"}, files=[ ("image", ("one.png", b"one", "image/png")), ("image", ("two.png", b"two", "image/png")), ], ) self.assertEqual(response.status_code, 200, response.text) self.assertEqual(response.json()["id"], "edit-1") self.assertEqual(len(self.fake_service.edit_calls), 1) self.assertEqual(self.fake_service.edit_calls[0][1]["account_pool_strategy"], "own_first") images = self.fake_service.edit_calls[0][1]["images"] self.assertEqual(len(images), 2) def test_create_edit_task_forwards_mask_upload(self): response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={"client_task_id": "edit-mask-1", "prompt": "outpaint", "model": "gpt-image-2"}, files=[ ("image", ("canvas.png", PNG_BYTES, "image/png")), ("mask", ("mask.png", PNG_BYTES, "image/png")), ], ) self.assertEqual(response.status_code, 200, response.text) payload = self.fake_service.edit_calls[-1][1] self.assertEqual(payload["masks"], [(PNG_BYTES, "mask.png", "image/png")]) def test_create_edit_task_preserves_memory_fields_from_multipart_form(self): response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={ "client_task_id": "edit-memory-1", "prompt": "continue this image conversation", "model": "gpt-image-2", "memory": "true", "memory_conversation_id": "local-conversation-1", "memory_reset": "false", "account_pool_strategy": "shared_first", }, files=[("image", ("one.png", b"one", "image/png"))], ) self.assertEqual(response.status_code, 200, response.text) payload = self.fake_service.edit_calls[-1][1] self.assertIs(payload["memory_enabled"], True) self.assertEqual(payload["memory_conversation_id"], "local-conversation-1") self.assertIs(payload["memory_reset"], False) self.assertEqual(payload["account_pool_strategy"], "shared_first") def test_create_edit_task_accepts_image_url(self): """测试图片编辑任务接口支持表单 image_url 引用。""" response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={ "client_task_id": "edit-url-1", "prompt": "edit", "model": "gpt-image-2", "image_url": DATA_IMAGE_URL, }, ) self.assertEqual(response.status_code, 200, response.text) self.assertEqual(len(self.fake_service.edit_calls), 1) images = self.fake_service.edit_calls[0][1]["images"] self.assertEqual(images, [(PNG_BYTES, "image_url.png", "image/png")]) def test_create_edit_task_can_use_third_party_image_api(self): fake_third_party = FakeThirdPartyImageService() with mock.patch.object(image_tasks_module, "third_party_image_service", fake_third_party): response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={ "client_task_id": "edit-api-1", "prompt": "edit", "model": "gpt-image-2", "image_url": DATA_IMAGE_URL, "quality": "high", "resolution": "1k", "third_party_image_group_id": "g1", }, ) self.assertEqual(response.status_code, 200, response.text) payload = self.fake_service.edit_calls[-1][1] self.assertEqual(payload["quality"], "high") self.assertEqual(payload["resolution"], "1k") self.assertEqual(payload["third_party_image_group_id"], "g1") self.assertEqual(fake_third_party.find_route_calls[0]["advance"], False) def test_named_free_api_model_without_group_routes_to_image_api_without_credit_cost(self): self.identity = {"id": "normal-1", "role": "normal", "name": "Normal", "account_pool_enabled": False} with mock.patch.object(image_tasks_module, "third_party_image_service", FakeFreeNamedThirdPartyImageService()): response = self.client.post( "/api/image-tasks/edits", headers=AUTH_HEADERS, data={ "client_task_id": "edit-free-1", "prompt": "edit", "model": "gpt-image-2（free）", "image_url": DATA_IMAGE_URL, }, ) self.assertEqual(response.status_code, 200, response.text) payload = self.fake_service.edit_calls[-1][1] self.assertEqual(payload["resolution"], "1k") self.assertEqual(payload["third_party_image_group_id"], "g-free") self.assertEqual(payload["credit_user_id"], "normal-1") self.assertEqual(payload["credit_amount"], 0) def test_list_tasks_reports_missing_ids(self): response = self.client.get("/api/image-tasks?ids=task-1,missing", headers=AUTH_HEADERS) self.assertEqual(response.status_code, 200, response.text) payload = response.json() self.assertEqual([item["id"] for item in payload["items"]], ["task-1"]) self.assertEqual(payload["missing_ids"], ["missing"]) if __name__ == "__main__": unittest.main()