102 lines
5.2 KiB
Python
102 lines
5.2 KiB
Python
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
|
|
from app.analytics_cache import AnalyticsSnapshotCache, SUMMARY_WINDOWS
|
|
from app.live import TrafficHistory
|
|
from app.store import AlertStore
|
|
|
|
|
|
class AnalyticsCacheTests(unittest.TestCase):
|
|
def test_archive_and_snapshot_workers_start_and_stop_independently(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
history = TrafficHistory("", retention_hours=24, max_events=1000, memory_events=1000)
|
|
cache = AnalyticsSnapshotCache(
|
|
store,
|
|
history,
|
|
threading.Event(),
|
|
interval_seconds=15,
|
|
archive_interval_seconds=1,
|
|
)
|
|
cache.start()
|
|
time.sleep(0.05)
|
|
self.assertTrue(cache._archive_thread.is_alive())
|
|
self.assertTrue(cache._snapshot_thread.is_alive())
|
|
cache.stop(timeout=0.5)
|
|
self.assertFalse(cache._archive_thread.is_alive())
|
|
self.assertFalse(cache._snapshot_thread.is_alive())
|
|
store.close()
|
|
|
|
def test_background_worker_has_no_dashboard_windows_until_requested(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
history = TrafficHistory("", retention_hours=24, max_events=1000, memory_events=1000)
|
|
cache = AnalyticsSnapshotCache(store, history, threading.Event(), interval_seconds=60)
|
|
self.assertEqual(cache._due_windows(), [])
|
|
cache.get(900)
|
|
self.assertEqual(cache._due_windows(), [900])
|
|
store.close()
|
|
|
|
def test_refresh_persists_all_dashboard_windows_in_sqlite_cache(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
history = TrafficHistory("", retention_hours=24, max_events=1000, memory_events=1000)
|
|
history.add({"id":"a","ts_ms":int(time.time()*1000),"timestamp":"x","type":"flow","src_ip":"10.0.0.2","dest_ip":"1.1.1.1","proto":"TCP","app_proto":"tls","direction":"outbound","bytes":50})
|
|
cache = AnalyticsSnapshotCache(store, history, threading.Event(), interval_seconds=60)
|
|
cache.refresh_all()
|
|
status = cache.status()
|
|
self.assertEqual({row["window_seconds"] for row in status["persisted"]}, set(SUMMARY_WINDOWS))
|
|
snapshot = cache.get(900)
|
|
self.assertEqual(snapshot["events"], 1)
|
|
self.assertEqual(snapshot["snapshot_source"], "sqlite-snapshot")
|
|
store.close()
|
|
|
|
def test_refresh_replaces_stale_sqlite_snapshot_with_current_analytics(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
history = TrafficHistory("", retention_hours=24, max_events=1000, memory_events=1000)
|
|
store.save_traffic_snapshot(900, {"events": 7, "timeline": [{"bucket": 1, "events": 7}]})
|
|
cache = AnalyticsSnapshotCache(store, history, threading.Event(), interval_seconds=60)
|
|
cache.refresh_all()
|
|
snapshot = cache.get(900)
|
|
self.assertIsNotNone(snapshot)
|
|
self.assertEqual(snapshot["events"], 0)
|
|
# The in-memory backend is test/dev-only, so it intentionally marks
|
|
# analytics incomplete. refresh_all must replace the stale snapshot.
|
|
self.assertFalse(snapshot["analytics_complete"])
|
|
self.assertEqual(snapshot["snapshot_source"], "sqlite-snapshot")
|
|
# Old SQLite traffic snapshots may exist after an upgrade, but they
|
|
# are no longer a data source for the dashboard.
|
|
self.assertEqual(store.traffic_snapshot(900)["events"], 0)
|
|
store.close()
|
|
|
|
def test_clear_traffic_snapshots_removes_persisted_windows(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
store.save_traffic_snapshot(900, {"events": 1})
|
|
store.save_traffic_snapshot(3600, {"events": 2})
|
|
self.assertEqual(store.clear_traffic_snapshots(), 2)
|
|
self.assertEqual(store.traffic_snapshot_status()["windows"], [])
|
|
store.close()
|
|
|
|
def test_arbitrary_five_hour_window_is_not_rounded_to_six_hours(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
store = AlertStore(os.path.join(td, "ids.db"))
|
|
history = TrafficHistory("", retention_hours=24, max_events=1000, memory_events=1000)
|
|
now = int(time.time() * 1000)
|
|
history.add({"id":"inside","ts_ms":now - 4 * 3600 * 1000,"timestamp":"x","type":"flow","direction":"outbound","src_ip":"10.0.0.2","dest_ip":"1.1.1.1","bytes":10})
|
|
history.add({"id":"outside","ts_ms":now - int(5.5 * 3600 * 1000),"timestamp":"x","type":"flow","direction":"outbound","src_ip":"10.0.0.2","dest_ip":"1.1.1.1","bytes":10})
|
|
cache = AnalyticsSnapshotCache(store, history, threading.Event(), interval_seconds=60)
|
|
cache.refresh_windows((18000,))
|
|
snapshot = cache.get(18000)
|
|
self.assertEqual(snapshot["window_seconds"], 18000)
|
|
self.assertEqual(snapshot["events"], 1)
|
|
store.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|