File size: 2,062 Bytes
15b8951 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import base64
import numpy as np
from fastapi.testclient import TestClient
from vlm.webconsole.qa_console.server import create_app
class FakeSource:
is_mock = True
is_connected = False
def get_latest_frame(self):
return np.zeros((48, 64, 3), dtype=np.uint8)
class FakeEngine:
loaded = True
def stream_infer(self, frame, prompt):
yield "I see "
yield "a phone [100, 100, 500, 500]."
def make_client():
app = create_app(FakeSource(), FakeEngine())
return TestClient(app)
def test_root_serves_html():
r = make_client().get("/")
assert r.status_code == 200
assert "<html" in r.text.lower()
def test_status_endpoint():
r = make_client().get("/status")
assert r.json() == {"connected": False, "mock": True}
def test_ws_streams_tokens_and_image():
client = make_client()
with client.websocket_connect("/ws") as ws:
ws.send_json({"prompt": "find the phone"})
msgs = []
while True:
m = ws.receive_json()
msgs.append(m)
if m["type"] in ("done", "error"):
break
types = [m["type"] for m in msgs]
assert "token" in types
assert types[-1] == "done"
tokens = "".join(m["text"] for m in msgs if m["type"] == "token")
assert "phone" in tokens
images = [m for m in msgs if m["type"] == "image"]
assert len(images) == 1
raw = base64.b64decode(images[0]["data"])
assert raw[:2] == b"\xff\xd8"
def test_ws_no_box_no_image():
class NoBoxEngine(FakeEngine):
def stream_infer(self, frame, prompt):
yield "nothing here"
app = create_app(FakeSource(), NoBoxEngine())
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"prompt": "anything"})
msgs = []
while True:
m = ws.receive_json()
msgs.append(m)
if m["type"] in ("done", "error"):
break
assert not any(m["type"] == "image" for m in msgs)
assert msgs[-1]["type"] == "done"
|