from __future__ import annotations import base64 import io from fastapi.testclient import TestClient from PIL import Image from app.main import app def png_bytes(color=(12, 80, 160)) -> bytes: buffer = io.BytesIO() Image.new("RGB", (64, 48), color).save(buffer, format="PNG") return buffer.getvalue() def data_url(data: bytes) -> str: return "data:image/png;base64," + base64.b64encode(data).decode() def test_health_ready(): client = TestClient(app) assert client.get("/health").json() == {"status": "ok"} assert client.get("/ready").json() == {"status": "ready"} def test_upload_then_followup_text_only_references_image(): client = TestClient(app) upload = client.post( "/v1/files", files={"file": ("design.png", png_bytes(), "image/png")}, data={"purpose": "vision"}, ) assert upload.status_code == 200 file_id = upload.json()["id"] first = client.post( "/v1/responses", json={ "model": "glm-5.2-visual-runtime", "input": [ { "role": "user", "content": [ {"type": "input_text", "text": "Remember the visual style of this image."}, {"type": "input_image", "file_id": file_id}, ], } ], }, ) assert first.status_code == 200 response_id = first.json()["id"] second = client.post( "/v1/responses", json={ "model": "glm-5.2-visual-runtime", "previous_response_id": response_id, "input": [{"role": "user", "content": [{"type": "input_text", "text": "Create slides using those colors."}]}], }, ) assert second.status_code == 200 body = second.json() assert body["conversation_id"] == first.json()["conversation_id"] assert "presentation" in body["output_text"].lower() def test_base64_image_and_chart_query_uses_chart_ir(): client = TestClient(app) response = client.post( "/v1/responses", json={ "model": "glm-5.2-visual-runtime", "input": [ { "role": "user", "content": [ {"type": "input_text", "text": "What was the Q3 value in this chart?"}, {"type": "input_image", "image_url": data_url(png_bytes((230, 230, 240)))}, ], } ], }, ) assert response.status_code == 200 assert "ChartIR" in response.json()["output_text"] assert "37" in response.json()["output_text"] def test_rejects_localhost_image_url(): client = TestClient(app) response = client.post( "/v1/responses", json={ "model": "glm-5.2-visual-runtime", "input": [ { "role": "user", "content": [ {"type": "input_text", "text": "Describe this image."}, {"type": "input_image", "image_url": "http://localhost/private.png"}, ], } ], }, ) assert response.status_code == 400 assert "private" in response.json()["detail"].lower()