|
| 1 | +"""Drives examples/example_api.py over socket.io, acting as the VolView client.""" |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import multiprocessing |
| 5 | +import os |
| 6 | +import socket |
| 7 | +import sys |
| 8 | +import threading |
| 9 | +import uuid |
| 10 | +from contextlib import contextmanager |
| 11 | + |
| 12 | +import numpy as np |
| 13 | +import pytest |
| 14 | +import socketio |
| 15 | +import uvicorn |
| 16 | +from aiohttp import web |
| 17 | + |
| 18 | +from volview_server.__main__ import create_app |
| 19 | + |
| 20 | +# The example's process pool must not fork from this multi-threaded process. |
| 21 | +multiprocessing.set_start_method("forkserver", force=True) |
| 22 | +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "examples")) |
| 23 | +from example_api import volview # noqa: E402 |
| 24 | + |
| 25 | +SIZE = 16 |
| 26 | +SOURCE = np.random.default_rng(0).integers(0, 255, SIZE**3, dtype=np.uint8) |
| 27 | + |
| 28 | + |
| 29 | +def vtk_image(values): |
| 30 | + return { |
| 31 | + "vtkClass": "vtkImageData", |
| 32 | + "direction": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], |
| 33 | + "extent": [0, SIZE - 1, 0, SIZE - 1, 0, SIZE - 1], |
| 34 | + "spacing": [1.0, 1.0, 1.0], |
| 35 | + "origin": [0.0, 0.0, 0.0], |
| 36 | + "pointData": { |
| 37 | + "vtkClass": "vtkDataSetAttributes", |
| 38 | + "activeScalars": 0, |
| 39 | + "arrays": [ |
| 40 | + { |
| 41 | + "data": { |
| 42 | + "vtkClass": "vtkDataArray", |
| 43 | + "size": values.size, |
| 44 | + "values": values.tobytes(), |
| 45 | + "dataType": "Uint8Array", |
| 46 | + "numberOfComponents": 1, |
| 47 | + "name": "Scalars", |
| 48 | + } |
| 49 | + } |
| 50 | + ], |
| 51 | + }, |
| 52 | + } |
| 53 | + |
| 54 | + |
| 55 | +def image_values(image): |
| 56 | + return np.frombuffer( |
| 57 | + image["pointData"]["arrays"][0]["data"]["values"], dtype=np.uint8 |
| 58 | + ) |
| 59 | + |
| 60 | + |
| 61 | +@contextmanager |
| 62 | +def background_loop(): |
| 63 | + loop = asyncio.new_event_loop() |
| 64 | + thread = threading.Thread(target=loop.run_forever) |
| 65 | + thread.start() |
| 66 | + try: |
| 67 | + yield loop |
| 68 | + finally: |
| 69 | + loop.call_soon_threadsafe(loop.stop) |
| 70 | + thread.join() |
| 71 | + loop.close() |
| 72 | + |
| 73 | + |
| 74 | +def listening_socket(): |
| 75 | + # Listening before the server starts, so clients never race its startup. |
| 76 | + sock = socket.create_server(("127.0.0.1", 0)) |
| 77 | + return sock, f"http://127.0.0.1:{sock.getsockname()[1]}" |
| 78 | + |
| 79 | + |
| 80 | +@contextmanager |
| 81 | +def serve_aiohttp(): |
| 82 | + sock, url = listening_socket() |
| 83 | + with background_loop() as loop: |
| 84 | + |
| 85 | + async def start(): |
| 86 | + runner = web.AppRunner(create_app(volview)) |
| 87 | + await runner.setup() |
| 88 | + await web.SockSite(runner, sock).start() |
| 89 | + return runner |
| 90 | + |
| 91 | + runner = asyncio.run_coroutine_threadsafe(start(), loop).result() |
| 92 | + try: |
| 93 | + yield url |
| 94 | + finally: |
| 95 | + asyncio.run_coroutine_threadsafe(runner.cleanup(), loop).result() |
| 96 | + |
| 97 | + |
| 98 | +async def not_found(scope, receive, send): |
| 99 | + await send({"type": "http.response.start", "status": 404, "headers": []}) |
| 100 | + await send({"type": "http.response.body", "body": b""}) |
| 101 | + |
| 102 | + |
| 103 | +@contextmanager |
| 104 | +def serve_asgi(): |
| 105 | + sock, url = listening_socket() |
| 106 | + server = uvicorn.Server(uvicorn.Config(volview(not_found), log_level="warning")) |
| 107 | + with background_loop() as loop: |
| 108 | + serving = asyncio.run_coroutine_threadsafe(server.serve([sock]), loop) |
| 109 | + try: |
| 110 | + yield url |
| 111 | + finally: |
| 112 | + server.should_exit = True |
| 113 | + serving.result() |
| 114 | + |
| 115 | + |
| 116 | +@pytest.fixture(params=[serve_aiohttp, serve_asgi], ids=["aiohttp", "asgi"]) |
| 117 | +def server_url(request): |
| 118 | + with request.param() as url: |
| 119 | + yield url |
| 120 | + |
| 121 | + |
| 122 | +class FakeVolViewClient: |
| 123 | + """Answers the server's client-store calls the way the VolView app would.""" |
| 124 | + |
| 125 | + def __init__(self, url): |
| 126 | + self.url = url |
| 127 | + self.sio = socketio.AsyncClient() |
| 128 | + self.results = {} |
| 129 | + self.streams = {} |
| 130 | + self.store_calls = [] |
| 131 | + self.images = {} |
| 132 | + self.sio.on("rpc:result", self._on_result) |
| 133 | + self.sio.on("stream:result", self._on_stream_result) |
| 134 | + self.sio.on("rpc:call", self._on_server_call) |
| 135 | + |
| 136 | + async def __aenter__(self): |
| 137 | + await self.sio.connect( |
| 138 | + f"{self.url}?clientId={uuid.uuid4().hex}", transports=["websocket"] |
| 139 | + ) |
| 140 | + return self |
| 141 | + |
| 142 | + async def __aexit__(self, *exc): |
| 143 | + await self.sio.disconnect() |
| 144 | + |
| 145 | + async def _on_result(self, data): |
| 146 | + self.results.pop(data["rpcId"]).set_result(data) |
| 147 | + |
| 148 | + async def _on_stream_result(self, data): |
| 149 | + await self.streams[data["rpcId"]].put(data) |
| 150 | + |
| 151 | + async def _on_server_call(self, data): |
| 152 | + store, prop_chain, args = data["args"] |
| 153 | + method = (store, ".".join(prop_chain)) |
| 154 | + self.store_calls.append(method) |
| 155 | + result = None |
| 156 | + if method == ("image-cache", "getVtkImageData"): |
| 157 | + result = vtk_image(SOURCE) |
| 158 | + elif method == ("images", "addVTKImageData"): |
| 159 | + self.images[args[0]] = image_values(args[1]) |
| 160 | + result = "blurred-id" |
| 161 | + elif method == ("image-cache", "updateVTKImageData"): |
| 162 | + self.images[args[0]] = image_values(args[1]) |
| 163 | + await self.sio.emit( |
| 164 | + "rpc:result", {"rpcId": data["rpcId"], "ok": True, "data": result} |
| 165 | + ) |
| 166 | + |
| 167 | + async def call(self, name, *args): |
| 168 | + rpc_id = uuid.uuid4().hex |
| 169 | + self.results[rpc_id] = asyncio.get_running_loop().create_future() |
| 170 | + await self.sio.emit( |
| 171 | + "rpc:call", {"rpcId": rpc_id, "name": name, "args": list(args)} |
| 172 | + ) |
| 173 | + return await self.results[rpc_id] |
| 174 | + |
| 175 | + async def stream(self, name, *args): |
| 176 | + rpc_id = uuid.uuid4().hex |
| 177 | + self.streams[rpc_id] = asyncio.Queue() |
| 178 | + await self.sio.emit( |
| 179 | + "stream:call", {"rpcId": rpc_id, "name": name, "args": list(args)} |
| 180 | + ) |
| 181 | + items = [] |
| 182 | + while True: |
| 183 | + result = await self.streams[rpc_id].get() |
| 184 | + if not result["ok"] or result["done"]: |
| 185 | + return items, result |
| 186 | + items.append(result["data"]) |
| 187 | + |
| 188 | + |
| 189 | +def run_client(url, scenario): |
| 190 | + async def main(): |
| 191 | + async with FakeVolViewClient(url) as client: |
| 192 | + return await scenario(client) |
| 193 | + |
| 194 | + return asyncio.run(main()) |
| 195 | + |
| 196 | + |
| 197 | +def test_rpc_returns_results(server_url): |
| 198 | + async def scenario(client): |
| 199 | + return await asyncio.gather(*(client.call("add", i, i) for i in range(20))) |
| 200 | + |
| 201 | + results = run_client(server_url, scenario) |
| 202 | + assert [(r["ok"], r["data"]) for r in results] == [(True, 2 * i) for i in range(20)] |
| 203 | + |
| 204 | + |
| 205 | +def test_rpc_failures_are_reported(server_url): |
| 206 | + async def scenario(client): |
| 207 | + return await client.call("does_not_exist"), await client.call("add", 1) |
| 208 | + |
| 209 | + unknown, bad_args = run_client(server_url, scenario) |
| 210 | + assert not unknown["ok"] and "does_not_exist" in unknown["error"] |
| 211 | + assert not bad_args["ok"] and "missing 1 required" in bad_args["error"] |
| 212 | + |
| 213 | + |
| 214 | +def test_stream_yields_every_item_then_done(server_url): |
| 215 | + items, final = run_client(server_url, lambda client: client.stream("progress")) |
| 216 | + assert [item["progress"] for item in items] == list(range(1, 101)) |
| 217 | + assert final["ok"] and final["done"] |
| 218 | + |
| 219 | + |
| 220 | +def test_median_filter_round_trips_through_client_stores(server_url): |
| 221 | + async def scenario(client): |
| 222 | + first = await client.call("medianFilter", "source-id", 1) |
| 223 | + first_calls, client.store_calls = client.store_calls, [] |
| 224 | + rerun = await client.call("medianFilter", "blurred-id", 1) |
| 225 | + return first, first_calls, rerun, client.store_calls, client.images |
| 226 | + |
| 227 | + first, first_calls, rerun, rerun_calls, images = run_client(server_url, scenario) |
| 228 | + |
| 229 | + assert first["ok"] and rerun["ok"] |
| 230 | + assert first_calls == [ |
| 231 | + ("image-cache", "getVtkImageData"), |
| 232 | + ("images", "addVTKImageData"), |
| 233 | + ("view", "setDataForAllViews"), |
| 234 | + ] |
| 235 | + assert rerun_calls == [ |
| 236 | + ("image-cache", "getVtkImageData"), |
| 237 | + ("image-cache", "updateVTKImageData"), |
| 238 | + ("view", "setDataForAllViews"), |
| 239 | + ] |
| 240 | + blurred = images["blurred-id"] |
| 241 | + assert blurred.size == SOURCE.size |
| 242 | + assert blurred.std() < SOURCE.std() |
0 commit comments