106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
import base64
|
||
import struct
|
||
import tempfile
|
||
import unittest
|
||
import uuid
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
|
||
from core.ark_client import ArkConfig
|
||
from tools.gpt_image import GptImageTool
|
||
|
||
|
||
def _png_stub(width: int = 1536, height: int = 864) -> bytes:
|
||
return b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR" + struct.pack(">II", width, height)
|
||
|
||
|
||
class _FakeArkClient:
|
||
json_call = None
|
||
|
||
def __init__(self, *_args, **_kwargs):
|
||
pass
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *_args):
|
||
pass
|
||
|
||
def post_json(self, endpoint, body, *, timeout_s=None):
|
||
type(self).json_call = (endpoint, body, timeout_s)
|
||
return {
|
||
"data": [{"b64_json": base64.b64encode(_png_stub()).decode("ascii")}],
|
||
"usage": {"output_tokens": 123},
|
||
}
|
||
|
||
|
||
class GptImageToolTests(unittest.TestCase):
|
||
def setUp(self):
|
||
_FakeArkClient.json_call = None
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.root = Path(self.tmp.name)
|
||
self.working_dir = self.root / "task"
|
||
self.working_dir.mkdir()
|
||
self.cfg = {
|
||
"model_id": "gpt-image-2",
|
||
"endpoint": "/images/generations",
|
||
"default_size": "auto",
|
||
"default_quality": "auto",
|
||
"request_timeout_s": 300,
|
||
"price_cny_per_image": 0,
|
||
}
|
||
self.tool = GptImageTool(
|
||
gw_cfg=ArkConfig(api_key="test", base_url="https://example.test/v1", raw={}),
|
||
image_variant_cfg=self.cfg,
|
||
variant_key="gpt_image",
|
||
working_dir=self.working_dir,
|
||
task_id=uuid.uuid4(),
|
||
user_id=uuid.uuid4(),
|
||
base_dir=self.working_dir,
|
||
user_root=self.root,
|
||
daily_limit=0,
|
||
)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _execute(self, **kwargs):
|
||
with (
|
||
patch("tools.gpt_image.ArkClient", _FakeArkClient),
|
||
patch("tools.gpt_image.quota_gate", return_value=""),
|
||
patch("tools.gpt_image.record_usage_safe"),
|
||
):
|
||
return self.tool.execute(**kwargs)
|
||
|
||
def test_text_to_image_forwards_size_and_quality(self):
|
||
result = self._execute(
|
||
prompt="draw a materials lab",
|
||
size="1536x864",
|
||
quality="high",
|
||
)
|
||
|
||
self.assertTrue(result.startswith("[gpt_image]"))
|
||
endpoint, body, _timeout = _FakeArkClient.json_call
|
||
self.assertEqual(endpoint, "/images/generations")
|
||
self.assertEqual(body["size"], "1536x864")
|
||
self.assertEqual(body["quality"], "high")
|
||
self.assertIn("size=1536x864", result)
|
||
self.assertIn("quality=high", result)
|
||
|
||
def test_size_validation(self):
|
||
self.assertEqual(self.tool._normalize_size("auto"), ("auto", ""))
|
||
self.assertEqual(self.tool._normalize_size("1536×864"), ("1536x864", ""))
|
||
for value in ("1000x1000", "4096x1024", "3072x512", "640x640", "bad"):
|
||
with self.subTest(value=value):
|
||
_size, error = self.tool._normalize_size(value)
|
||
self.assertTrue(error.startswith("[Error]"))
|
||
|
||
def test_quality_validation(self):
|
||
result = self._execute(prompt="draw", quality="ultra")
|
||
self.assertEqual(result, "[Error] quality 必须是 auto / low / medium / high")
|
||
self.assertIsNone(_FakeArkClient.json_call)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|