| 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 |
|
|
