tabras / tests /test_art.py
vvennelakanti's picture
finished base
a8afc36
Raw
History Blame Contribute Delete
3.9 kB
from io import BytesIO
from PIL import Image
from art import (
DiffusersImageClient,
best_torch_dtype,
card_art_prompt,
illustrate_card,
illustrate_cards,
image_data_uri,
move_pipe_to_device,
theme_style,
)
from budget import Card
from primitives import Effect
class FakeArtClient:
# Return a deterministic image URI for prompts.
def create_art(self, prompt: str) -> str:
return f"uri:{prompt}"
class BrokenArtClient:
# Raise like a failing image backend.
def create_art(self, prompt: str) -> str:
raise RuntimeError("no art")
class FakePipe:
# Capture pipeline calls and return a PIL image.
def __init__(self) -> None:
self.calls = []
self.device = ""
def __call__(self, **kwargs):
self.calls.append(kwargs)
return type("Result", (), {"images": [Image.new("RGB", (2, 2), "red")]})()
def to(self, device: str):
self.device = device
return self
class FakeTorch:
float16 = "f16"
float32 = "f32"
class cuda:
# Report CUDA as unavailable.
@staticmethod
def is_available() -> bool:
return False
class backends:
class mps:
# Report MPS as available.
@staticmethod
def is_available() -> bool:
return True
# Verify images are encoded as browser-renderable PNG data URIs.
def test_image_data_uri() -> None:
uri = image_data_uri(Image.new("RGB", (1, 1), "blue"))
assert uri.startswith("data:image/png;base64,")
# Verify diffusers image client calls the pipeline and returns encoded art.
def test_diffusers_image_client_create_art() -> None:
pipe = FakePipe()
client = DiffusersImageClient(pipe, steps=1, width=64, height=64)
uri = client.create_art("dark spell")
assert uri.startswith("data:image/png;base64,")
assert pipe.calls[0]["prompt"] == "dark spell"
assert pipe.calls[0]["num_inference_steps"] == 1
assert pipe.calls[0]["width"] == 64
assert "text" in pipe.calls[0]["negative_prompt"]
# Verify dtype and device helpers prefer MPS when CUDA is unavailable.
def test_torch_helpers() -> None:
pipe = FakePipe()
assert best_torch_dtype(FakeTorch) == "f16"
assert move_pipe_to_device(pipe, FakeTorch) is pipe
assert pipe.device == "mps"
# Verify card prompts use model-authored art prompts first.
def test_card_art_prompt() -> None:
card = Card("Spark", 1, "fire", "wuxia", (Effect("deal", amount=2),), art_prompt="ember saint")
assert "ember saint" in card_art_prompt(card)
assert "no text" in card_art_prompt(card)
fallback = Card("Spark", 1, "fire", "wuxia", (Effect("deal", amount=2),))
prompt = card_art_prompt(fallback)
assert "Spark" in prompt
assert "Deal 2 damage." in prompt
assert "wide wuxia fantasy spell illustration" in prompt
# Verify user worlds become concrete visual style clauses.
def test_theme_style() -> None:
assert "anime fantasy key visual" in theme_style("Anime")
assert "dark fantasy spell illustration" in theme_style("dark fantasy")
# Verify illustration attaches generated art without changing card rules.
def test_illustrate_card() -> None:
card = Card("Spark", 1, "fire", "wuxia", (Effect("deal", amount=2),), art_prompt="ember")
illustrated = illustrate_card(FakeArtClient(), card)
assert illustrated.art_uri.startswith("uri:wide wuxia fantasy spell illustration")
assert illustrated.rules_text() == card.rules_text()
assert illustrate_card(None, card) is card
assert illustrate_card(BrokenArtClient(), card) is card
# Verify illustration maps across packs.
def test_illustrate_cards() -> None:
cards = (Card("Spark", 1, "fire", "wuxia", (Effect("deal", amount=2),), art_prompt="ember"),)
assert illustrate_cards(FakeArtClient(), cards)[0].art_uri.startswith("uri:wide wuxia")