| """CPU regression checks for image forwarding, probabilities and input errors.""" |
| import base64 |
| import io |
| import json |
| import unittest |
|
|
| import httpx |
| from PIL import Image |
|
|
| import decisions_server as server |
| from decisions_api import ApiError, assemble, build_plan |
| from image_inputs import image_input |
| from image_jevbench import OpenJevImageAdapter |
|
|
|
|
| def picture(): |
| b = io.BytesIO() |
| Image.new("RGB", (32, 32), "red").save(b, format="PNG") |
| return b.getvalue() |
|
|
|
|
| class Images(unittest.IsolatedAsyncioTestCase): |
| async def asyncSetUp(self): |
| self.calls = [] |
| def upstream(req): |
| body = json.loads(req.content) |
| self.calls.append(body) |
| return httpx.Response(200, json=[{"embedding": [0, 2, -1], "meta_info": {"prompt_tokens": 80}} |
| for _ in body["text"]]) |
| server._client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) |
| self.api = httpx.AsyncClient(transport=httpx.ASGITransport(app=server.app), base_url="http://test") |
| server.API_KEY = "" |
| self.body = {"model": "test", "state": "Image: <<IMG>>", "image_data": base64.b64encode(picture()).decode(), |
| "questions": {"q": {"type": "choice", "instructions": "Which color?", |
| "criteria": {"a": "red", "b": "blue"}}}} |
|
|
| async def asyncTearDown(self): |
| await self.api.aclose() |
| await server._client.aclose() |
| server._client = None |
|
|
| async def test_pixels_forwarded_per_option(self): |
| r = await self.api.post("/v1/systemone", json=self.body) |
| self.assertEqual(r.status_code, 200, r.text) |
| self.assertEqual(r.json()["answers"]["q"]["probabilities"], {"a": .5, "b": .5}) |
| self.assertEqual(len(self.calls), 1) |
| for c in self.calls: |
| self.assertEqual(c["image_data"], [self.body["image_data"]] * 2) |
| self.assertEqual(c["text"][0].count("<|image_pad|>"), 1) |
| self.assertNotIn("<<IMG>>", c["text"][0]) |
|
|
| async def test_text_request_still_batched(self): |
| del self.body["image_data"] |
| self.body["state"] = "A red box." |
| r = await self.api.post("/v1/systemone", json=self.body) |
| self.assertEqual(r.status_code, 200) |
| self.assertEqual(len(self.calls), 1) |
| self.assertNotIn("image_data", self.calls[0]) |
|
|
| async def test_bad_image_does_not_reach_worker(self): |
| for data in ("/etc/passwd", "https://example.com/image.png", "not-base64", [], ""): |
| self.body["image_data"] = data |
| r = await self.api.post("/v1/systemone", json=self.body) |
| self.assertEqual(r.status_code, 400, r.text) |
| self.assertEqual(self.calls, []) |
|
|
| async def test_duplicate_marker_rejected(self): |
| self.body["state"] = "<<IMG>> <<IMG>>" |
| r = await self.api.post("/v1/systemone", json=self.body) |
| self.assertEqual(r.status_code, 400) |
| self.assertEqual(self.calls, []) |
|
|
| def test_data_uri_and_zero_entailment(self): |
| self.assertEqual(image_input({"image_data": "data:image/png;base64," + self.body["image_data"]}), self.body["image_data"]) |
| p = assemble(build_plan(self.body), [0, 0])["q"]["probabilities"] |
| self.assertEqual(p, {"a": .5, "b": .5}) |
| with self.assertRaises(ValueError): |
| assemble(build_plan(self.body), [float("nan"), 1]) |
|
|
| def test_benchmark_does_not_send_gold_or_alt(self): |
| a = OpenJevImageAdapter("http://test", "test") |
| try: |
| body = a.build_request({"question": "Q", "options": [{"label": "a", "text": "red"}], |
| "correctLabel": "SECRET_GOLD", "alt": "SECRET_DESCRIPTION"}, picture()) |
| self.assertNotIn("SECRET", json.dumps(body)) |
| finally: |
| a.close() |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
