64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import gpt_image2_generate as image_demo
|
|
import gpt_image2_generate_multi as image_multi_demo
|
|
|
|
|
|
class ImageDemoHelpersTest(unittest.TestCase):
|
|
def test_openai_ratio_size_mapping(self) -> None:
|
|
self.assertEqual(image_demo.resolve_api_size("16:9", "openai"), "1536x864")
|
|
self.assertEqual(image_demo.resolve_api_size("3840x2160", "openai"), "3840x2160")
|
|
|
|
def test_min_output_edge(self) -> None:
|
|
self.assertEqual(image_demo.resolve_min_output_edge("4096x2304", None, "apicodex"), 4096)
|
|
self.assertIsNone(image_demo.resolve_min_output_edge("1024x1024", None, "openai"))
|
|
|
|
|
|
class ImageMultiDemoHelpersTest(unittest.TestCase):
|
|
def test_collect_reference_sources_from_json_array(self) -> None:
|
|
with patch.dict(os.environ, {"REFERENCE_IMAGES": '["./a.png", "./b.png", "./a.png"]'}, clear=True):
|
|
self.assertEqual(image_multi_demo.collect_reference_sources(), ["./a.png", "./b.png"])
|
|
|
|
def test_build_edit_files_field_names(self) -> None:
|
|
ref1 = image_demo.ImageInput("a.png", "image/png", b"a")
|
|
ref2 = image_demo.ImageInput("b.png", "image/png", b"b")
|
|
|
|
openai_single = image_multi_demo.build_edit_files("openai", [ref1], None)
|
|
openai_multi = image_multi_demo.build_edit_files("openai", [ref1, ref2], None)
|
|
apicodex_multi = image_multi_demo.build_edit_files("apicodex", [ref1, ref2], None)
|
|
|
|
self.assertEqual([name for name, _ in openai_single], ["image"])
|
|
self.assertEqual([name for name, _ in openai_multi], ["image[]", "image[]"])
|
|
self.assertEqual([name for name, _ in apicodex_multi], ["image[]", "image[]"])
|
|
|
|
def test_output_stem_env_is_sanitized(self) -> None:
|
|
with patch.dict(os.environ, {"IMAGE_OUTPUT_STEM": "../bad stem!"}, clear=True):
|
|
self.assertEqual(image_multi_demo.resolve_output_stem(), "bad_stem")
|
|
|
|
def test_default_n_matches_reference_count(self) -> None:
|
|
env = {
|
|
"IMAGE_PROMPT": "hello",
|
|
"IMAGE_SIZE": "1:1",
|
|
}
|
|
with patch.dict(os.environ, env, clear=True):
|
|
payload, model, size, request_mode, mask_source, background, input_fidelity = image_multi_demo.build_payload(
|
|
"http://localhost",
|
|
["./a.png", "./b.png"],
|
|
)
|
|
|
|
self.assertEqual(payload["n"], 2)
|
|
self.assertEqual(request_mode, "edit")
|
|
self.assertIsNone(mask_source)
|
|
self.assertIsNone(background)
|
|
self.assertIsNone(input_fidelity)
|
|
self.assertEqual(model, image_demo.DEFAULT_MODEL)
|
|
self.assertEqual(size, "1:1")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|