import io, time
from PIL import Image
from fastapi.testclient import TestClient
from app import app

c = TestClient(app)
colors = ["red","green","blue","yellow","purple"]
def png(col):
    b = io.BytesIO(); Image.new("RGB", (640, 480), col).save(b, "PNG"); return b.getvalue()

r = c.post("/generate", data={"prompt": "What color is this image?", "max_new_tokens": 20},
           files=[("images", ("a.png", png("red"), "image/png"))])
print("1 img:", r.status_code, r.json())

t = time.time()
files = [("images", (f"{i}.png", png(colors[i % 5]), "image/png")) for i in range(20)]
r = c.post("/generate", data={"prompt": "List the colors of these images.", "max_new_tokens": 60}, files=files)
print("20 img:", r.status_code, r.json(), f"{time.time()-t:.1f}s")
assert r.status_code == 200

files.append(("images", ("x.png", png("red"), "image/png")))
r = c.post("/generate", data={"prompt": "hi"}, files=files)
print("21 img:", r.status_code); assert r.status_code == 400

r = c.post("/generate", data={"prompt": "hi"}, files=[("images", ("bad.png", b"notimage", "image/png"))])
print("bad img:", r.status_code); assert r.status_code == 400

r = c.post("/generate", data={"prompt": "Say hello.", "max_new_tokens": 10})
print("text only:", r.status_code, r.json()); assert r.status_code == 200