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