69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
import base64
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from polza import PolzaClient, PolzaError
|
|
|
|
|
|
def test_reference_payload_accepts_url_and_data_uri():
|
|
assert PolzaClient.reference_payload("https://example.com/style.png") == {
|
|
"type": "url",
|
|
"data": "https://example.com/style.png",
|
|
}
|
|
data_uri = "data:image/png;base64,AAAA"
|
|
assert PolzaClient.reference_payload(data_uri) == {"type": "base64", "data": data_uri}
|
|
|
|
|
|
def test_reference_payload_encodes_local_image(tmp_path):
|
|
path = tmp_path / "style.png"
|
|
path.write_bytes(b"image-bytes")
|
|
|
|
payload = PolzaClient.reference_payload(str(path))
|
|
|
|
assert payload["type"] == "base64"
|
|
assert payload["data"].startswith("data:image/png;base64,")
|
|
assert base64.b64decode(payload["data"].split(",", 1)[1]) == b"image-bytes"
|
|
|
|
|
|
def test_reference_payload_rejects_unknown_path():
|
|
with pytest.raises(PolzaError, match="does not exist"):
|
|
PolzaClient.reference_payload("missing-style.png")
|
|
|
|
|
|
def test_create_image_sends_references_and_variants(monkeypatch):
|
|
client = PolzaClient(api_key="test-key")
|
|
captured = {}
|
|
|
|
def fake_request(method, path, payload=None):
|
|
captured.update(method=method, path=path, payload=payload)
|
|
return {"id": "gen_1", "status": "pending"}
|
|
|
|
monkeypatch.setattr(client, "_request", fake_request)
|
|
client.create_image(
|
|
model="seedream-3",
|
|
prompt="game sprite",
|
|
reference_images=["https://example.com/style.png"],
|
|
count=3,
|
|
aspect_ratio="1:1",
|
|
seed=42,
|
|
wait=False,
|
|
)
|
|
|
|
assert captured["method"] == "POST"
|
|
assert captured["path"] == "/media"
|
|
assert captured["payload"]["async"] is True
|
|
assert captured["payload"]["input"]["max_images"] == 3
|
|
assert captured["payload"]["input"]["images"][0]["type"] == "url"
|
|
|
|
|
|
def test_image_sources_handles_urls_and_base64():
|
|
response = {
|
|
"data": [{"url": "https://cdn.example/one.png"}],
|
|
"result": {"images": [{"b64_json": "aGVsbG8="}]},
|
|
}
|
|
assert PolzaClient.image_sources(response) == [
|
|
"https://cdn.example/one.png",
|
|
"aGVsbG8=",
|
|
]
|