| 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() |
|
|
