205 lines
6.9 KiB
Python
205 lines
6.9 KiB
Python
import importlib
|
|
import json
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
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()
|
|
async def load_current(keys):
|
|
assert keys == ["1"]
|
|
return [{"mmsi": 1, "lat": 60.0, "lon": 11.0, "source": "barentswatch_vessels"}]
|
|
|
|
monkeypatch.setattr(broadcaster, "_load_current_vessel_updates", load_current)
|
|
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": 60.0,
|
|
"lon": 11.0,
|
|
"source": "barentswatch_vessels",
|
|
"action": "upsert",
|
|
"created": None,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_global_vessel_subscription_delivers_every_item_in_bounded_frames():
|
|
manager = ConnectionManager()
|
|
socket = FakeWebSocket()
|
|
config = manager.subscribe_vessels(socket, {"scope": "global", "zoom": 4})
|
|
json.dumps(config)
|
|
vessels = [{"mmsi": index, "lat": 60, "lon": 10} for index in range(2501)]
|
|
await manager.broadcast_vessels({"vessels": vessels})
|
|
|
|
assert [len(frame["payload"]["vessels"]) for frame in socket.sent] == [1000, 1000, 501]
|
|
assert [item for frame in socket.sent for item in frame["payload"]["vessels"]] == vessels
|
|
manager.unsubscribe(socket, ["vessels"])
|
|
await manager.broadcast_vessels({"vessels": vessels})
|
|
assert len(socket.sent) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_global_vessel_removal_does_not_require_coordinates():
|
|
manager = ConnectionManager()
|
|
socket = FakeWebSocket()
|
|
manager.subscribe_vessels(socket, {"scope": "global", "zoom": 4})
|
|
await manager.broadcast_vessels({"vessels": [{"mmsi": 123, "action": "remove"}]})
|
|
assert socket.sent[0]["payload"]["vessels"] == [{"mmsi": 123, "action": "remove"}]
|
|
|
|
|
|
def test_anonymous_earth_can_confirm_global_vessel_subscription(monkeypatch):
|
|
websocket_module = importlib.import_module("app.api.v1.websocket")
|
|
monkeypatch.setattr(websocket_module, "manager", ConnectionManager())
|
|
app = FastAPI()
|
|
app.include_router(websocket_module.router)
|
|
with TestClient(app) as client, client.websocket_connect("/ws") as socket:
|
|
assert socket.receive_json()["type"] == "connection_established"
|
|
socket.send_json({
|
|
"type": "subscribe",
|
|
"data": {"channels": ["earth_updates", "vessels"], "scope": "global", "zoom": 4},
|
|
})
|
|
response = socket.receive_json()
|
|
assert response["type"] == "subscription_confirmed"
|
|
assert response["data"]["vessels"]["scope"] == "global"
|
|
assert response["data"]["vessels"]["type"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vessel_flush_retries_without_overwriting_newer_queued_updates(monkeypatch):
|
|
broadcaster = DataBroadcaster()
|
|
broadcaster.enqueue_vessel_update({"vessels": [{"mmsi": 1, "lat": 59, "lon": 10}]})
|
|
|
|
async def fail_read(_keys):
|
|
broadcaster.enqueue_vessel_update({"vessels": [{"mmsi": 1, "lat": 61, "lon": 12}]})
|
|
raise RuntimeError("database unavailable")
|
|
|
|
monkeypatch.setattr(broadcaster, "_load_current_vessel_updates", fail_read)
|
|
with pytest.raises(RuntimeError, match="database unavailable"):
|
|
await broadcaster.flush_vessel_updates()
|
|
assert broadcaster._pending_vessel_updates["1"]["lat"] == 61
|