151 lines
4.7 KiB
Python
151 lines
4.7 KiB
Python
import pytest
|
|
|
|
from app.core.websocket.manager import ConnectionManager
|
|
from app.core.websocket.ue_scene import (
|
|
UeSceneStateStore,
|
|
_item_revision,
|
|
build_incremental_changes,
|
|
)
|
|
|
|
|
|
class _FakeSession:
|
|
pass
|
|
|
|
|
|
class _FakeWebSocket:
|
|
def __init__(self):
|
|
self.accepted = False
|
|
self.messages = []
|
|
|
|
async def accept(self):
|
|
self.accepted = True
|
|
|
|
async def send_json(self, message):
|
|
self.messages.append(message)
|
|
|
|
async def close(self):
|
|
return None
|
|
|
|
|
|
def _scene_item(item_id: str, title: str, *, lat: float = 0.0, lng: float = 0.0):
|
|
item = {
|
|
"id": item_id,
|
|
"entity_type": "gpu_cluster",
|
|
"geo": {"lat": lat, "lng": lng, "alt": 0.0},
|
|
"visual": {"style": "pulse_marker", "size": 1.0, "color": "#FF8C42"},
|
|
"metrics": {"name": title},
|
|
"labels": {"title": title, "subtitle": ""},
|
|
"status": {"health": "normal", "alert_level": "none"},
|
|
}
|
|
item["revision"] = _item_revision(item)
|
|
return item
|
|
|
|
|
|
def _scene_state(state_hash: int, items_by_layer: dict[str, dict[str, dict]]):
|
|
base_layers = {
|
|
"satellites": {"revision": 0, "items": {}},
|
|
"supercomputers": {"revision": 0, "items": {}},
|
|
"gpu_clusters": {"revision": 0, "items": {}},
|
|
"submarine_cables": {"revision": 0, "items": {}},
|
|
"landing_points": {"revision": 0, "items": {}},
|
|
"alerts": {"revision": 0, "items": {}},
|
|
}
|
|
for layer_name, items in items_by_layer.items():
|
|
base_layers[layer_name]["items"] = items
|
|
base_layers[layer_name]["revision"] = len(items)
|
|
return {
|
|
"generated_at": "2026-04-17T00:00:00Z",
|
|
"state_hash": state_hash,
|
|
"total_records": sum(len(items) for items in items_by_layer.values()),
|
|
"layers": base_layers,
|
|
}
|
|
|
|
|
|
def test_build_incremental_changes_detects_add_update_remove():
|
|
previous_state = _scene_state(
|
|
1,
|
|
{
|
|
"gpu_clusters": {
|
|
"gpu:a": _scene_item("gpu:a", "A"),
|
|
"gpu:b": _scene_item("gpu:b", "B"),
|
|
}
|
|
},
|
|
)
|
|
current_state = _scene_state(
|
|
2,
|
|
{
|
|
"gpu_clusters": {
|
|
"gpu:b": _scene_item("gpu:b", "B Updated"),
|
|
"gpu:c": _scene_item("gpu:c", "C"),
|
|
}
|
|
},
|
|
)
|
|
|
|
changes = build_incremental_changes(previous_state, current_state)
|
|
|
|
assert changes["gpu_clusters"]["added"][0]["id"] == "gpu:c"
|
|
assert changes["gpu_clusters"]["updated"][0]["id"] == "gpu:b"
|
|
assert changes["gpu_clusters"]["removed"] == ["gpu:a"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_state_store_replays_incremental_history(monkeypatch):
|
|
state_list = [
|
|
_scene_state(1, {"gpu_clusters": {"gpu:a": _scene_item("gpu:a", "A")}}),
|
|
_scene_state(
|
|
2,
|
|
{"gpu_clusters": {"gpu:a": _scene_item("gpu:a", "A"), "gpu:b": _scene_item("gpu:b", "B")}},
|
|
),
|
|
_scene_state(
|
|
3,
|
|
{"gpu_clusters": {"gpu:a": _scene_item("gpu:a", "A Updated"), "gpu:b": _scene_item("gpu:b", "B")}},
|
|
),
|
|
]
|
|
state_index = {"value": 0}
|
|
|
|
async def _fake_build_visualization_scene_state(_db):
|
|
current_index = state_index["value"]
|
|
if current_index < len(state_list) - 1:
|
|
state_index["value"] += 1
|
|
return state_list[current_index]
|
|
|
|
monkeypatch.setattr(
|
|
"app.core.websocket.ue_scene.build_visualization_scene_state",
|
|
_fake_build_visualization_scene_state,
|
|
)
|
|
|
|
store = UeSceneStateStore(history_limit=5)
|
|
session = _FakeSession()
|
|
|
|
first_payloads = await store.get_broadcast_payloads(session)
|
|
second_payloads = await store.get_broadcast_payloads(session)
|
|
third_payloads = await store.get_broadcast_payloads(session)
|
|
replay_payloads = await store.get_sync_payloads(
|
|
session,
|
|
last_sequence=1,
|
|
reason="sequence_gap",
|
|
)
|
|
|
|
assert first_payloads[0]["update_type"] == "full"
|
|
assert second_payloads[0]["update_type"] == "incremental"
|
|
assert third_payloads[0]["sequence"] == 3
|
|
assert [payload["sequence"] for payload in replay_payloads] == [2, 3]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connection_manager_broadcasts_to_subscribed_channel_only():
|
|
manager = ConnectionManager()
|
|
ws_dashboard = _FakeWebSocket()
|
|
ws_scene = _FakeWebSocket()
|
|
|
|
await manager.connect(ws_dashboard, "user-dashboard")
|
|
await manager.connect(ws_scene, "user-scene")
|
|
|
|
manager.subscribe(ws_dashboard, ["dashboard"])
|
|
manager.subscribe(ws_scene, ["ue_scene"])
|
|
|
|
await manager.broadcast({"type": "data_frame", "channel": "ue_scene"}, channel="ue_scene")
|
|
|
|
assert ws_dashboard.messages == []
|
|
assert ws_scene.messages == [{"type": "data_frame", "channel": "ue_scene"}]
|