289 lines
10 KiB
Python
289 lines
10 KiB
Python
import importlib.util
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import call, patch
|
|
|
|
MODULE_PATH = Path(__file__).parents[1] / "generate.py"
|
|
SPEC = importlib.util.spec_from_file_location("logo_generate", MODULE_PATH)
|
|
logo_generate = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(logo_generate)
|
|
|
|
|
|
class AtlasGenerationTests(unittest.TestCase):
|
|
@patch.object(logo_generate, "_download_atlas_image")
|
|
@patch.object(logo_generate.time, "sleep")
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_atlas_submits_once_and_polls_until_completed(
|
|
self, json_request, sleep, download
|
|
):
|
|
json_request.side_effect = [
|
|
{"code": 200, "data": {"id": "pred-123", "status": "created"}},
|
|
{"code": 200, "data": {"id": "pred-123", "status": "processing"}},
|
|
{
|
|
"code": 200,
|
|
"data": {
|
|
"id": "pred-123",
|
|
"status": "completed",
|
|
"outputs": ["https://media.example.com/logo.png"],
|
|
},
|
|
},
|
|
]
|
|
|
|
logo_generate._generate_with_atlas(
|
|
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
|
)
|
|
|
|
self.assertEqual(json_request.call_count, 3)
|
|
self.assertEqual(
|
|
json_request.call_args_list[0],
|
|
call(
|
|
f"{logo_generate.ATLAS_API_BASE}/model/generateImage",
|
|
"atlas-key",
|
|
method="POST",
|
|
payload={
|
|
"model": "atlas/model",
|
|
"prompt": "logo prompt",
|
|
"aspect_ratio": "1:1",
|
|
},
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
json_request.call_args_list[1:],
|
|
[
|
|
call(
|
|
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
|
"atlas-key",
|
|
),
|
|
call(
|
|
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
|
"atlas-key",
|
|
),
|
|
],
|
|
)
|
|
self.assertEqual(sleep.call_count, 2)
|
|
download.assert_called_once_with(
|
|
"https://media.example.com/logo.png", "logo.png"
|
|
)
|
|
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_atlas_does_not_retry_generation_post(self, json_request):
|
|
json_request.side_effect = RuntimeError("network error")
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "network error"):
|
|
logo_generate._generate_with_atlas(
|
|
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
|
)
|
|
|
|
json_request.assert_called_once()
|
|
|
|
@patch.object(logo_generate, "_validate_public_https_url")
|
|
@patch.object(logo_generate, "build_opener")
|
|
def test_media_download_never_forwards_api_key(self, build_opener, validate):
|
|
class Headers:
|
|
@staticmethod
|
|
def get_content_type():
|
|
return "image/png"
|
|
|
|
class Response:
|
|
headers = Headers()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return None
|
|
|
|
@staticmethod
|
|
def read():
|
|
return b"png-bytes"
|
|
|
|
build_opener.return_value.open.return_value = Response()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
output = Path(temp_dir) / "logo.png"
|
|
logo_generate._download_atlas_image(
|
|
"https://media.example.com/logo.png", output
|
|
)
|
|
self.assertEqual(output.read_bytes(), b"png-bytes")
|
|
|
|
request = build_opener.return_value.open.call_args.args[0]
|
|
headers = {key.lower(): value for key, value in request.header_items()}
|
|
self.assertNotIn("authorization", headers)
|
|
self.assertEqual(headers["accept"], "image/*")
|
|
self.assertEqual(headers["user-agent"], logo_generate.HTTP_USER_AGENT)
|
|
validate.assert_called_once_with("https://media.example.com/logo.png")
|
|
|
|
def test_atlas_requires_api_key(self):
|
|
with self.assertRaisesRegex(RuntimeError, "ATLASCLOUD_API_KEY not set"):
|
|
logo_generate._generate_with_atlas(
|
|
"logo prompt", "logo.png", "1:1", None, "atlas/model"
|
|
)
|
|
|
|
def test_media_url_rejects_private_addresses(self):
|
|
with self.assertRaisesRegex(ValueError, "non-public address"):
|
|
logo_generate._validate_public_https_url("https://127.0.0.1/logo.png")
|
|
|
|
with self.assertRaisesRegex(ValueError, "local hostname"):
|
|
logo_generate._validate_public_https_url("https://assets.local/logo.png")
|
|
|
|
|
|
class MuapiGenerationTests(unittest.TestCase):
|
|
@patch.object(logo_generate, "_download_muapi_image")
|
|
@patch.object(logo_generate.time, "sleep")
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_muapi_submits_once_and_polls_until_completed(
|
|
self, json_request, sleep, download
|
|
):
|
|
json_request.side_effect = [
|
|
{
|
|
"id": "req-123",
|
|
"status": "created",
|
|
"output": {
|
|
"urls": {
|
|
"get": "https://api.muapi.ai/api/v1/results/req-123"
|
|
}
|
|
},
|
|
},
|
|
{"id": "req-123", "status": "processing"},
|
|
{
|
|
"id": "req-123",
|
|
"status": "completed",
|
|
"output": {"outputs": ["https://media.example.com/logo.png"]},
|
|
},
|
|
]
|
|
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "nano-banana"
|
|
)
|
|
|
|
self.assertEqual(json_request.call_count, 3)
|
|
self.assertEqual(
|
|
json_request.call_args_list[0],
|
|
call(
|
|
f"{logo_generate.MUAPI_API_BASE}/nano-banana",
|
|
"muapi-key",
|
|
method="POST",
|
|
payload={"prompt": "logo prompt", "aspect_ratio": "1:1"},
|
|
api_key_header="x-api-key",
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
json_request.call_args_list[1:],
|
|
[
|
|
call(
|
|
"https://api.muapi.ai/api/v1/results/req-123",
|
|
"muapi-key",
|
|
api_key_header="x-api-key",
|
|
),
|
|
call(
|
|
"https://api.muapi.ai/api/v1/results/req-123",
|
|
"muapi-key",
|
|
api_key_header="x-api-key",
|
|
),
|
|
],
|
|
)
|
|
self.assertEqual(sleep.call_count, 2)
|
|
download.assert_called_once_with(
|
|
"https://media.example.com/logo.png", "logo.png"
|
|
)
|
|
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_muapi_does_not_retry_generation_post(self, json_request):
|
|
json_request.side_effect = RuntimeError("network error")
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "network error"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "nano-banana"
|
|
)
|
|
|
|
json_request.assert_called_once()
|
|
|
|
def test_muapi_requires_key_and_known_model(self):
|
|
with self.assertRaisesRegex(RuntimeError, "MUAPI_API_KEY not set"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", None, "nano-banana"
|
|
)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Unsupported MuAPI logo model"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "unknown-model"
|
|
)
|
|
|
|
@patch.object(logo_generate, "build_opener")
|
|
def test_muapi_uses_x_api_key_header(self, build_opener):
|
|
class Response:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return None
|
|
|
|
@staticmethod
|
|
def read():
|
|
return b"{}"
|
|
|
|
build_opener.return_value.open.return_value = Response()
|
|
|
|
logo_generate._json_request(
|
|
"https://api.muapi.ai/api/v1/nano-banana",
|
|
"muapi-key",
|
|
method="POST",
|
|
payload={"prompt": "logo"},
|
|
api_key_header="x-api-key",
|
|
)
|
|
|
|
request = build_opener.return_value.open.call_args.args[0]
|
|
headers = {key.lower(): value for key, value in request.header_items()}
|
|
self.assertEqual(headers["x-api-key"], "muapi-key")
|
|
self.assertNotIn("authorization", headers)
|
|
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_muapi_reports_failed_prediction(self, json_request):
|
|
json_request.side_effect = [
|
|
{
|
|
"request_id": "req-123",
|
|
"output": {
|
|
"urls": {
|
|
"get": "https://api.muapi.ai/api/v1/results/req-123"
|
|
}
|
|
},
|
|
},
|
|
{"status": "failed", "error": "invalid prompt"},
|
|
]
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "invalid prompt"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "nano-banana"
|
|
)
|
|
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_muapi_requires_creation_result_url(self, json_request):
|
|
json_request.return_value = {"request_id": "req-123", "status": "created"}
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "valid HTTPS result URL"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "nano-banana"
|
|
)
|
|
|
|
json_request.assert_called_once()
|
|
|
|
@patch.object(logo_generate, "_json_request")
|
|
def test_muapi_rejects_invalid_creation_result_url(self, json_request):
|
|
json_request.return_value = {
|
|
"request_id": "req-123",
|
|
"status": "created",
|
|
"output": {"urls": {"get": "http://api.muapi.ai/results/req-123"}},
|
|
}
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "valid HTTPS result URL"):
|
|
logo_generate._generate_with_muapi(
|
|
"logo prompt", "logo.png", "1:1", "muapi-key", "nano-banana"
|
|
)
|
|
|
|
json_request.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|