"""Offline tests for the Seedream image transport.""" import sys import unittest from pathlib import Path from unittest.mock import AsyncMock, patch from PIL import Image ROOT = Path(__file__).resolve().parents[1] CUSTOM_NODES_ROOT = ROOT.parent COMFY_ROOT = CUSTOM_NODES_ROOT.parent sys.path.insert(0, str(COMFY_ROOT)) sys.path.insert(0, str(CUSTOM_NODES_ROOT)) from comfyui_o1key.clients import seedream_image_client as CLIENT # noqa: E402 class SeedreamBodyTests(unittest.TestCase): def test_model_routes_resolve_to_the_documented_identifier(self): for route in ("畅速", "直连", "专线"): self.assertEqual( CLIENT.resolve_seedream_model("Seedream 5.0 Pro", route), CLIENT.SEEDREAM_API_MODEL_ID, ) def test_body_keeps_reference_order_and_fixed_watermark(self): body = CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt=" 一只宇航猫 ", size="2368x1776", output_format="PNG", image_urls=[ "https://example.invalid/one.png", "https://example.invalid/two.jpg", ], ) self.assertEqual(body, { "model": CLIENT.SEEDREAM_API_MODEL_ID, "prompt": "一只宇航猫", "n": 1, "size": "2368x1776", "output_format": "png", "watermark": False, "images": [ "https://example.invalid/one.png", "https://example.invalid/two.jpg", ], }) def test_body_rejects_webp_and_non_https_references(self): with self.assertRaisesRegex(ValueError, "png 或 jpeg"): CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="test", size="1024x1024", output_format="webp", ) with self.assertRaisesRegex(ValueError, "HTTPS URL"): CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="test", size="1024x1024", output_format="jpeg", image_urls=["http://example.invalid/reference.png"], ) def test_body_rejects_sizes_outside_the_documented_matrix(self): with self.assertRaisesRegex(ValueError, "图片尺寸无效"): CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="test", size="1920x1080", output_format="jpeg", ) def test_body_omits_size_for_smart_resolution(self): body = CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="test", size=None, output_format="jpeg", ) self.assertNotIn("size", body) def test_reference_dimensions_match_current_volcengine_limits(self): CLIENT.validate_seedream_reference_dimensions(15, 15) CLIENT.validate_seedream_reference_dimensions(240, 15) CLIENT.validate_seedream_reference_dimensions(15, 240) with self.assertRaisesRegex(ValueError, "宽和高都必须大于 14px"): CLIENT.validate_seedream_reference_dimensions(14, 240) with self.assertRaisesRegex(ValueError, "宽高比必须在 1:16～16:1"): CLIENT.validate_seedream_reference_dimensions(241, 15) with self.assertRaisesRegex(ValueError, "总像素不能超过"): CLIENT.validate_seedream_reference_dimensions(6001, 6000) def test_layer_reference_uses_its_documented_total_pixel_floor(self): CLIENT.validate_seedream_reference_dimensions( 512, 512, layer_decomposition=True, ) with self.assertRaisesRegex(ValueError, "总像素必须在 512×512"): CLIENT.validate_seedream_reference_dimensions( 511, 512, layer_decomposition=True, ) def test_reference_file_size_is_checked_against_exact_upload_payload(self): class OversizedPayload: def __len__(self): return CLIENT.SEEDREAM_REFERENCE_MAX_BYTES + 1 source = Image.new("RGB", (15, 15), "red") with ( patch.object( CLIENT, "image_to_upload_payload", return_value=(OversizedPayload(), ".png", "image/png"), ), self.assertRaisesRegex(ValueError, "文件不能超过 30MB"), ): CLIENT.validate_seedream_reference_image(source) def test_layer_decomposition_allows_blank_prompt_and_requires_one_png_reference(self): body = CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="", size="1.5K", output_format="png", image_urls=["https://example.invalid/poster.png"], layer_decomposition=True, ) self.assertNotIn("prompt", body) self.assertEqual(body["size"], "1.5K") self.assertIs(body["layer_decomposition"], True) with self.assertRaisesRegex(ValueError, "必须且只能提供1张"): CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="", size="auto", output_format="png", layer_decomposition=True, ) with self.assertRaisesRegex(ValueError, "仅支持 png"): CLIENT.build_seedream_submit_body( model=CLIENT.SEEDREAM_API_MODEL_ID, prompt="", size="2K", output_format="jpeg", image_urls=["https://example.invalid/poster.png"], layer_decomposition=True, ) def test_layer_metadata_is_sanitized_without_result_urls(self): metadata = CLIENT.extract_seedream_layer_metadata({ "data": { "images": [{ "url": "https://signed.invalid/secret.png", "z_index": "1", "size": "1408x1780", "output_format": "png", "bounding_box": { "absolute": [320, 180, 1728, 1960], "normalized": [156, 88, 844, 957], }, "name": "主体", "description": "画面中的主要人物", }], }, }) self.assertEqual(metadata[0]["z_index"], 1) self.assertEqual(metadata[0]["bounding_box"]["absolute"], [320, 180, 1728, 1960]) self.assertNotIn("url", metadata[0]) class SeedreamLifecycleTests(unittest.IsolatedAsyncioTestCase): async def test_upload_submit_poll_and_parse_use_the_shared_async_lifecycle(self): source = Image.new("RGB", (15, 15), "red") result = Image.new("RGB", (2, 2), "blue") completed = { "task_id": "task-seedream", "status": "SUCCESS", "data": {"images": [{"url": "https://example.invalid/result.png"}]}, } metrics = { "download_bytes": 123, "download_seconds": 0.2, "download_wall_seconds": 0.1, "inline_images": 0, } with ( patch.object( CLIENT, "upload_images_to_temp_urls", new=AsyncMock(return_value=["https://example.invalid/reference.png"]), ) as upload, patch.object( CLIENT, "submit_async_image_task", new=AsyncMock(return_value="task-seedream"), ) as submit, patch.object( CLIENT, "poll_async_image_task", new=AsyncMock(return_value=completed), ) as poll, patch.object( CLIENT, "parse_completed_async_image_task", new=AsyncMock(return_value=(completed, ([result], metrics))), ) as parse, ): client = CLIENT.SeedreamImageClient( base_url="https://cf-api.o1key.com/", api_key="secret", ) images, timing = await client.generate_async( session=object(), prompt="test", model=CLIENT.SEEDREAM_API_MODEL_ID, size="1024x1024", output_format="jpeg", images=[source], ) self.assertEqual(images, [result]) upload.assert_awaited_once() self.assertEqual(upload.await_args.kwargs["base_url"], "https://cf-api.o1key.com") body = submit.await_args.args[3] self.assertEqual(body["images"], ["https://example.invalid/reference.png"]) self.assertEqual(body["output_format"], "jpeg") self.assertIs(body["watermark"], False) poll.assert_awaited_once() parse.assert_awaited_once() self.assertEqual(timing["task_id"], "task-seedream") self.assertEqual(timing["download_bytes"], 123) async def test_layer_results_are_sorted_by_z_index_and_keep_safe_metadata(self): top = Image.new("RGBA", (2, 2), (0, 0, 255, 64)) base = Image.new("RGB", (2, 2), "white") completed = { "status": "SUCCESS", "data": {"images": [ {"url": "https://signed.invalid/top.png", "z_index": 2, "name": "文字"}, {"url": "https://signed.invalid/base.png", "z_index": 0, "name": "底图"}, ]}, } with ( patch.object(CLIENT, "upload_images_to_temp_urls", new=AsyncMock( return_value=["https://example.invalid/reference.png"] )), patch.object(CLIENT, "submit_async_image_task", new=AsyncMock(return_value="task-layer")), patch.object(CLIENT, "poll_async_image_task", new=AsyncMock(return_value=completed)), patch.object(CLIENT, "parse_completed_async_image_task", new=AsyncMock( return_value=(completed, ([top, base], { "download_bytes": 1, "download_seconds": 0.0, "download_wall_seconds": 0.0, "inline_images": 0, })) )), ): images, timing = await CLIENT.SeedreamImageClient( base_url="https://cf-api.o1key.com", api_key="secret", ).generate_async( session=object(), prompt="", model=CLIENT.SEEDREAM_API_MODEL_ID, size="auto", output_format="png", images=[Image.new("RGB", (512, 512), "red")], layer_decomposition=True, ) self.assertEqual([item["z_index"] for item in timing["result_metadata"]], [0, 2]) self.assertEqual(getattr(images[0], "_o1key_seedream_layer")["name"], "底图") async def test_invalid_reference_is_rejected_before_upload(self): upload = AsyncMock(return_value=["https://example.invalid/reference.png"]) with patch.object(CLIENT, "upload_images_to_temp_urls", new=upload): with self.assertRaisesRegex(ValueError, "宽和高都必须大于 14px"): await CLIENT.SeedreamImageClient( base_url="https://cf.invalid", api_key="secret", ).generate_async( session=object(), prompt="test", model=CLIENT.SEEDREAM_API_MODEL_ID, size="1024x1024", output_format="jpeg", images=[Image.new("RGB", (14, 15), "red")], ) upload.assert_not_awaited() if __name__ == "__main__": unittest.main(verbosity=2)