stroblme
8c50c22444
ws:// and wss:// URLs dial out and hold a supervised connection; a path is served under /hooks/<flow>/ with an optional secret. Outputs only listen, inputs send. Frames are JSON keyed by port, with json_key and on_connect.
137 lines
4.6 KiB
Python
137 lines
4.6 KiB
Python
"""The websocket node, both roles: served under /hooks, or dialled out."""
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
from websockets.asyncio.server import ServerConnection, serve
|
|
|
|
from fluksio.flow.controller import HOOK_PREFIX
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.nodes import WebSocketNode
|
|
from fluksio.flow.pipeline import Pipeline
|
|
|
|
SECRET = "s0cket-key"
|
|
HOOK = f"{HOOK_PREFIX}/house/ticker"
|
|
|
|
|
|
def ws_node(url: str, *, sends: bool = False, **params: object) -> WebSocketNode:
|
|
port = [MessageSpec(name="temp", dtype=DType.FLOAT)]
|
|
node = WebSocketNode(
|
|
requires=port if sends else (),
|
|
provides=() if sends else port,
|
|
params={"url": url, **params},
|
|
name="ticker",
|
|
)
|
|
node.assign_flow("house", "ticker")
|
|
if node.serves:
|
|
node.url = HOOK # what the controller does
|
|
return node
|
|
|
|
|
|
def test_frames_map_to_ports():
|
|
node = ws_node("ws://x", json_key="value")
|
|
assert node._decode('{"value": "21.5"}') == {"temp": 21.5}
|
|
# One output takes a frame that names no port whole.
|
|
assert node._decode("19") == {"temp": 19.0}
|
|
|
|
sender = ws_node("ws://x", sends=True)
|
|
assert sender._encode({"temp": 21.5}) == "21.5"
|
|
assert sender._encode({"cmd": {"op": "on"}}) == '{"op": "on"}'
|
|
assert json.loads(sender._encode({"a": 1, "b": "x"})) == {"a": 1, "b": "x"}
|
|
|
|
with pytest.raises(ValueError):
|
|
ws_node("http://x")
|
|
|
|
|
|
def test_a_served_socket_injects_frames_and_guards_its_secret():
|
|
node = ws_node("/ticker", secret=SECRET)
|
|
pipeline = Pipeline(nodes=[node])
|
|
app = FastAPI()
|
|
app.mount("/", FastAPI()) # the MCP catch-all the route must sit ahead of
|
|
node.register_route(app)
|
|
client = TestClient(app)
|
|
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect(f"{HOOK}/wrong") as ws:
|
|
ws.receive_text()
|
|
|
|
with client.websocket_connect(f"{HOOK}/{SECRET}") as ws:
|
|
ws.send_text('{"temp": 21.5}')
|
|
assert pipeline.values()["house.temp"]["value"] == 21.5
|
|
|
|
|
|
def test_a_served_sender_reaches_every_connected_client():
|
|
node = ws_node("/ticker", sends=True)
|
|
app = FastAPI()
|
|
|
|
with TestClient(app) as client:
|
|
client.portal.call(node.start, app)
|
|
with (
|
|
client.websocket_connect(HOOK) as first,
|
|
client.websocket_connect(HOOK) as second,
|
|
):
|
|
# Both are registered once either has been accepted and spoken to.
|
|
client.portal.call(asyncio.sleep, 0.05)
|
|
node._sender_handler({}, temp=21.5)
|
|
assert first.receive_text() == "21.5"
|
|
assert second.receive_text() == "21.5"
|
|
client.portal.call(node.stop, app)
|
|
|
|
|
|
def test_a_client_listens_and_sends_over_a_held_connection():
|
|
received: list[str] = []
|
|
|
|
async def server(conn: ServerConnection) -> None:
|
|
if conn.request is not None and conn.request.path == "/feed":
|
|
await conn.send('{"temp": 19.5}')
|
|
await conn.wait_closed()
|
|
else:
|
|
async for frame in conn:
|
|
received.append(str(frame))
|
|
|
|
async def scenario() -> None:
|
|
async with serve(server, "127.0.0.1", 0) as srv:
|
|
port = srv.sockets[0].getsockname()[1]
|
|
listener = ws_node(f"ws://127.0.0.1:{port}/feed")
|
|
pipeline = Pipeline(nodes=[listener])
|
|
sender = ws_node(
|
|
f"ws://127.0.0.1:{port}/cmd",
|
|
sends=True,
|
|
on_connect=[{"op": "subscribe"}],
|
|
)
|
|
await listener.start()
|
|
await sender.start()
|
|
await asyncio.to_thread(sender._sender_handler, {}, temp=21.5)
|
|
for _ in range(100):
|
|
if len(received) == 2 and pipeline.values():
|
|
break
|
|
await asyncio.sleep(0.02)
|
|
await listener.stop()
|
|
await sender.stop()
|
|
|
|
assert received == ['{"op": "subscribe"}', "21.5"]
|
|
assert pipeline.values()["house.temp"]["value"] == 19.5
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_start_respawns_a_loop_cancelled_underneath_it():
|
|
"""Enabling a flow cancels its supervised loops and starts the nodes again,
|
|
without a stop in between."""
|
|
node = ws_node("/ticker", sends=True)
|
|
|
|
async def scenario() -> None:
|
|
await node.start()
|
|
first = node._task
|
|
assert first is not None
|
|
await node._cancel_task(first)
|
|
await node.start()
|
|
assert node._task is not first and not node._task.done()
|
|
await node.stop()
|
|
|
|
asyncio.run(scenario())
|