139 lines
4.0 KiB
Python
139 lines
4.0 KiB
Python
import pytest
|
|
import importlib
|
|
|
|
from app.core.websocket.manager import ConnectionManager
|
|
from app.core.websocket.broadcaster import DataBroadcaster
|
|
|
|
|
|
class FakeWebSocket:
|
|
def __init__(self):
|
|
self.accepted = False
|
|
self.sent = []
|
|
self.closed = False
|
|
|
|
async def accept(self):
|
|
self.accepted = True
|
|
|
|
async def send_json(self, message):
|
|
self.sent.append(message)
|
|
|
|
async def close(self):
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_channel_subscribers_receive_channel_broadcasts():
|
|
manager = ConnectionManager()
|
|
socket = FakeWebSocket()
|
|
|
|
await manager.connect(socket, "user-1")
|
|
manager.subscribe(socket, ["dashboard"])
|
|
await manager.broadcast({"type": "data_frame", "channel": "dashboard"}, channel="dashboard")
|
|
|
|
assert socket.accepted is True
|
|
assert socket.sent == [{"type": "data_frame", "channel": "dashboard"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_removes_channel_subscriptions():
|
|
manager = ConnectionManager()
|
|
socket = FakeWebSocket()
|
|
|
|
await manager.connect(socket, "user-1")
|
|
manager.subscribe(socket, ["dashboard"])
|
|
manager.disconnect(socket, "user-1")
|
|
await manager.broadcast({"type": "data_frame", "channel": "dashboard"}, channel="dashboard")
|
|
|
|
assert socket.sent == []
|
|
assert "dashboard" not in manager.channel_subscriptions
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vessel_subscribers_receive_only_matching_bbox_updates():
|
|
manager = ConnectionManager()
|
|
oslo_socket = FakeWebSocket()
|
|
bergen_socket = FakeWebSocket()
|
|
|
|
await manager.connect(oslo_socket, "user-1")
|
|
await manager.connect(bergen_socket, "user-2")
|
|
manager.subscribe_vessels(
|
|
oslo_socket,
|
|
{"bbox": [10, 59, 11, 60], "zoom": 12, "limit": 1000},
|
|
)
|
|
manager.subscribe_vessels(
|
|
bergen_socket,
|
|
{"bbox": [5, 60, 6, 61], "zoom": 12, "limit": 1000},
|
|
)
|
|
|
|
await manager.broadcast_vessels(
|
|
{
|
|
"action": "upsert",
|
|
"vessels": [
|
|
{"mmsi": 1, "lat": 59.9, "lon": 10.7},
|
|
{"mmsi": 2, "lat": 60.3, "lon": 5.3},
|
|
],
|
|
}
|
|
)
|
|
|
|
assert oslo_socket.sent[0]["payload"]["vessels"] == [{"mmsi": 1, "lat": 59.9, "lon": 10.7}]
|
|
assert bergen_socket.sent[0]["payload"]["vessels"] == [{"mmsi": 2, "lat": 60.3, "lon": 5.3}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vessel_broadcast_removes_slow_connections():
|
|
manager = ConnectionManager()
|
|
|
|
class BrokenWebSocket(FakeWebSocket):
|
|
async def send_json(self, message):
|
|
raise RuntimeError("client is gone")
|
|
|
|
socket = BrokenWebSocket()
|
|
await manager.connect(socket, "user-1")
|
|
manager.subscribe_vessels(socket, {"bbox": [10, 59, 11, 60], "zoom": 12})
|
|
|
|
await manager.broadcast_vessels({"vessels": [{"mmsi": 1, "lat": 59.9, "lon": 10.7}]})
|
|
|
|
assert socket not in manager.vessel_subscriptions
|
|
|
|
|
|
def test_vessel_subscription_rejects_large_bbox():
|
|
manager = ConnectionManager()
|
|
|
|
with pytest.raises(ValueError, match="bbox is too large"):
|
|
manager.subscribe_vessels(FakeWebSocket(), {"bbox": [-180, -90, 180, 90], "zoom": 1})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vessel_broadcaster_keeps_latest_update_per_mmsi(monkeypatch):
|
|
sent = []
|
|
|
|
async def fake_broadcast_vessels(payload):
|
|
sent.append(payload)
|
|
|
|
broadcaster_module = importlib.import_module("app.core.websocket.broadcaster")
|
|
monkeypatch.setattr(broadcaster_module.manager, "broadcast_vessels", fake_broadcast_vessels)
|
|
broadcaster = DataBroadcaster()
|
|
broadcaster.enqueue_vessel_update(
|
|
{
|
|
"source": "aisstream_vessels",
|
|
"vessels": [
|
|
{"mmsi": 1, "lat": 59.0, "lon": 10.0},
|
|
{"mmsi": 1, "lat": 59.1, "lon": 10.1},
|
|
],
|
|
}
|
|
)
|
|
|
|
await broadcaster.flush_vessel_updates()
|
|
|
|
assert len(sent) == 1
|
|
assert sent[0]["vessels"] == [
|
|
{
|
|
"mmsi": 1,
|
|
"lat": 59.1,
|
|
"lon": 10.1,
|
|
"source": "aisstream_vessels",
|
|
"action": "upsert",
|
|
"created": None,
|
|
}
|
|
]
|