From 5c947e6c2e3fdf89880954a5a1f369b0f861e1fa Mon Sep 17 00:00:00 2001 From: Lloyd Date: Wed, 22 Apr 2026 09:48:27 +0100 Subject: [PATCH] feat: enhance WebSocket handling and add throttling for stats broadcasting --- .../data_acquisition/storage_collector.py | 37 ++++++--- .../data_acquisition/websocket_handler.py | 5 ++ ...est_storage_collector_ws_stats_throttle.py | 79 +++++++++++++++++++ 3 files changed, 110 insertions(+), 11 deletions(-) create mode 100644 tests/test_storage_collector_ws_stats_throttle.py diff --git a/repeater/data_acquisition/storage_collector.py b/repeater/data_acquisition/storage_collector.py index 8fa5652..ec92406 100644 --- a/repeater/data_acquisition/storage_collector.py +++ b/repeater/data_acquisition/storage_collector.py @@ -60,11 +60,19 @@ class StorageCollector: # Initialize WebSocket handler for real-time updates self.websocket_available = False + self.websocket_has_connected_clients = lambda: False + self._last_ws_stats_broadcast: float = 0.0 + self._ws_stats_broadcast_interval_sec: float = 5.0 try: - from .websocket_handler import broadcast_packet, broadcast_stats + from .websocket_handler import ( + broadcast_packet, + broadcast_stats, + has_connected_clients, + ) self.websocket_broadcast_packet = broadcast_packet self.websocket_broadcast_stats = broadcast_stats + self.websocket_has_connected_clients = has_connected_clients self.websocket_available = True logger.info("WebSocket handler initialized for real-time updates") except ImportError: @@ -182,16 +190,23 @@ class StorageCollector: if self.websocket_available: try: self.websocket_broadcast_packet(packet_record) - packet_stats_24h = self.sqlite_handler.get_packet_stats(hours=24) - uptime_seconds = ( - time.time() - self.repeater_handler.start_time if self.repeater_handler else 0 - ) - self.websocket_broadcast_stats( - { - "packet_stats": packet_stats_24h, - "system_stats": {"uptime_seconds": uptime_seconds}, - } - ) + if self.websocket_has_connected_clients(): + now_mono = time.monotonic() + if ( + now_mono - self._last_ws_stats_broadcast + >= self._ws_stats_broadcast_interval_sec + ): + self._last_ws_stats_broadcast = now_mono + packet_stats_24h = self.sqlite_handler.get_packet_stats(hours=24) + uptime_seconds = ( + time.time() - self.repeater_handler.start_time if self.repeater_handler else 0 + ) + self.websocket_broadcast_stats( + { + "packet_stats": packet_stats_24h, + "system_stats": {"uptime_seconds": uptime_seconds}, + } + ) except Exception as e: logger.debug(f"WebSocket broadcast failed: {e}") diff --git a/repeater/data_acquisition/websocket_handler.py b/repeater/data_acquisition/websocket_handler.py index 2e87ebc..bf88b66 100644 --- a/repeater/data_acquisition/websocket_handler.py +++ b/repeater/data_acquisition/websocket_handler.py @@ -126,6 +126,11 @@ def broadcast_stats(stats_data: dict): _connected_clients.discard(client) +def has_connected_clients() -> bool: + """Return True when at least one authenticated websocket client is connected.""" + return bool(_connected_clients) + + def _heartbeat_loop(): """Background thread to send periodic pings to all connected clients""" global _heartbeat_running diff --git a/tests/test_storage_collector_ws_stats_throttle.py b/tests/test_storage_collector_ws_stats_throttle.py new file mode 100644 index 0000000..d6e4814 --- /dev/null +++ b/tests/test_storage_collector_ws_stats_throttle.py @@ -0,0 +1,79 @@ +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +sys.modules.setdefault("psutil", types.ModuleType("psutil")) + +nacl_module = types.ModuleType("nacl") +nacl_signing_module = types.ModuleType("nacl.signing") + + +class _SigningKeyStub: + pass + + +nacl_signing_module.SigningKey = _SigningKeyStub +nacl_module.signing = nacl_signing_module + +sys.modules.setdefault("nacl", nacl_module) +sys.modules.setdefault("nacl.signing", nacl_signing_module) + +from repeater.data_acquisition.storage_collector import StorageCollector + + +def _make_collector() -> StorageCollector: + with ( + patch("repeater.data_acquisition.storage_collector.SQLiteHandler"), + patch("repeater.data_acquisition.storage_collector.RRDToolHandler"), + patch("repeater.data_acquisition.hardware_stats.HardwareStatsCollector"), + ): + collector = StorageCollector(config={"storage": {"storage_dir": "/tmp/pymc_repeater_test"}}) + + collector.sqlite_handler = MagicMock() + collector.sqlite_handler.get_packet_stats.return_value = {"total_packets": 1} + collector.websocket_available = True + collector.websocket_broadcast_packet = MagicMock() + collector.websocket_broadcast_stats = MagicMock() + collector.websocket_has_connected_clients = MagicMock(return_value=True) + collector.repeater_handler = SimpleNamespace(start_time=100.0) + return collector + + +def test_publish_packet_sync_first_call_broadcasts_stats_immediately(): + collector = _make_collector() + + with patch("repeater.data_acquisition.storage_collector.time.monotonic", return_value=1000.0): + collector._publish_packet_sync({"type": 1, "transmitted": True}, skip_mqtt=False) + + assert collector.sqlite_handler.get_packet_stats.call_count == 1 + assert collector.websocket_broadcast_stats.call_count == 1 + assert collector.websocket_broadcast_packet.call_count == 1 + + +def test_publish_packet_sync_throttles_stats_to_interval(): + collector = _make_collector() + call_times = [1000.0 + (i * 0.1) for i in range(10)] + [1005.1] + + with patch( + "repeater.data_acquisition.storage_collector.time.monotonic", + side_effect=call_times, + ): + for _ in call_times: + collector._publish_packet_sync({"type": 1, "transmitted": True}, skip_mqtt=False) + + assert collector.sqlite_handler.get_packet_stats.call_count == 2 + assert collector.websocket_broadcast_stats.call_count == 2 + assert collector.websocket_broadcast_packet.call_count == len(call_times) + + +def test_publish_packet_sync_always_broadcasts_packet_event_even_without_clients(): + collector = _make_collector() + collector.websocket_has_connected_clients.return_value = False + + for _ in range(5): + collector._publish_packet_sync({"type": 1, "transmitted": True}, skip_mqtt=False) + + assert collector.websocket_broadcast_packet.call_count == 5 + assert collector.sqlite_handler.get_packet_stats.call_count == 0 + assert collector.websocket_broadcast_stats.call_count == 0