forked from whudqw/pinai-imagegen
73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
import pathlib
|
|
import sys
|
|
import unittest
|
|
|
|
SCRIPT_DIR = pathlib.Path(__file__).resolve().parents[1] / "scripts"
|
|
SKILL_DIR = SCRIPT_DIR.parent
|
|
sys.path.insert(0, str(SCRIPT_DIR))
|
|
|
|
try:
|
|
from generate_image import DEFAULT_API_URL, extract_image_bytes, images_endpoint, png_metadata
|
|
except ModuleNotFoundError:
|
|
DEFAULT_API_URL = None
|
|
extract_image_bytes = None
|
|
images_endpoint = None
|
|
png_metadata = None
|
|
|
|
|
|
class ExtractImageBytesTests(unittest.TestCase):
|
|
def test_decodes_inline_base64_image(self):
|
|
self.assertIsNotNone(extract_image_bytes, "generate_image.extract_image_bytes is required")
|
|
payload = {"data": [{"b64_json": "aGVsbG8="}]}
|
|
self.assertEqual(extract_image_bytes(payload), b"hello")
|
|
|
|
def test_downloads_url_response_through_injected_reader(self):
|
|
self.assertIsNotNone(extract_image_bytes, "generate_image.extract_image_bytes is required")
|
|
payload = {"data": [{"url": "https://images.example.test/output.png"}]}
|
|
self.assertEqual(extract_image_bytes(payload, download_url=lambda url: b"png-bytes"), b"png-bytes")
|
|
|
|
def test_rejects_unknown_image_payload(self):
|
|
self.assertIsNotNone(extract_image_bytes, "generate_image.extract_image_bytes is required")
|
|
with self.assertRaisesRegex(ValueError, "b64_json or url"):
|
|
extract_image_bytes({"data": [{}]})
|
|
|
|
|
|
class ImageMetadataTests(unittest.TestCase):
|
|
def test_uses_pinai_as_the_default_api_url(self):
|
|
self.assertEqual(DEFAULT_API_URL, "https://api.pinaic.com/v1")
|
|
|
|
def test_uses_a_base_url_or_complete_endpoint(self):
|
|
self.assertIsNotNone(images_endpoint, "generate_image.images_endpoint is required")
|
|
self.assertEqual(
|
|
images_endpoint("https://provider.example/v1"),
|
|
"https://provider.example/v1/images/generations",
|
|
)
|
|
self.assertEqual(
|
|
images_endpoint("https://provider.example/v1/images/generations/"),
|
|
"https://provider.example/v1/images/generations",
|
|
)
|
|
|
|
def test_reads_png_dimensions_and_color_mode(self):
|
|
self.assertIsNotNone(png_metadata, "generate_image.png_metadata is required")
|
|
png = b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR" + b"\x00\x00\x04\x00\x00\x00\x02\x00\x08\x06\x00\x00\x00"
|
|
self.assertEqual(png_metadata(png), (1024, 512, 8, "RGBA"))
|
|
|
|
|
|
class KeySetupPackageTests(unittest.TestCase):
|
|
def test_includes_a_double_click_key_setup_launcher(self):
|
|
launcher = SKILL_DIR / "设置 PinAI Key.cmd"
|
|
setup_script = SCRIPT_DIR / "setup_pinai_key.ps1"
|
|
self.assertTrue(launcher.is_file())
|
|
self.assertTrue(setup_script.is_file())
|
|
self.assertIn("setup_pinai_key.ps1", launcher.read_text(encoding="utf-8"))
|
|
|
|
def test_key_setup_uses_a_masked_local_input_and_user_variable(self):
|
|
setup_script = (SCRIPT_DIR / "setup_pinai_key.ps1").read_text(encoding="utf-8")
|
|
self.assertIn("UseSystemPasswordChar = $true", setup_script)
|
|
self.assertIn("SetEnvironmentVariable('PINAI_API_KEY', $key, 'User')", setup_script)
|
|
self.assertNotIn("Write-Output $key", setup_script)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|