Files
vacuum-wall/tests/test_ws_snapshot.py
T
mteehan 332d14e37d ws: migrate push stream to data streaming
- daemon: send full snapshot on connect; versions/tick now carry the
  full state of one subsystem (subsystem + data); no legacy
  updated/subsystems payloads; refresh_state and POST /status/refresh
  broadcast per-subsystem versions with data
- client: modelSet() patches models in place; onMessage/topic refresh
  retired; 3s initial-load fallback via new POST /api/status/refresh
- schema: lib/schema.py TypedDicts + hoover/schema.js defaults +
  docs/state-model.md as single source of truth for state shapes
- system: poll at 1s, volatile metrics registered, dashboard uses a
  dedicated system model (status model removed)
- firewall: refuse to strip both https and ssh from the default zone
  (409, force override via UI confirm); set_zone_services persists
  services to the declarative config; collector exposes default_zone
- UI: pages migrate to flat state shapes; post-mutation modelFetch
  refreshes removed (WS delta covers it)
- tests: ws snapshot/delta/broadcast, refresh-state, schema types,
  model-set/js ws handler and reconnect fallback
2026-08-20 01:38:00 +00:00

128 lines
4.1 KiB
Python

"""Tests for the WS connect snapshot (daemon.server._handle_ws).
After the push-stream migration, a successful WS handshake sends a full
state snapshot ({type: snapshot, data: {subsystem: state|null, ...}})
instead of the retired {type: init, versions: ...} message.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from lib.state import State
@pytest.fixture(autouse=False)
def db_reset():
"""Isolated in-memory DB so a builtin admin exists for token minting.
Mirrors the autouse _db_reset fixture in tests/test_auth.py (the DB
singleton must be reset and pointed at SQLite :memory: before each test).
"""
import os
from lib.db import get_db, reset_db_for_test
reset_db_for_test()
old_backend = os.environ.pop("VACUUM_WALL_DB_BACKEND", None)
old_path = os.environ.pop("VACUUM_WALL_DB_PATH", None)
old_seed = os.environ.pop("VACUUM_WALL_SEED_BUILTIN_ADMIN", None)
os.environ["VACUUM_WALL_DB_BACKEND"] = "sqlite"
os.environ["VACUUM_WALL_DB_PATH"] = ":memory:"
get_db() # triggers builtin-admin seed on the empty :memory: DB
yield
reset_db_for_test()
if old_backend is not None:
os.environ["VACUUM_WALL_DB_BACKEND"] = old_backend
if old_path is not None:
os.environ["VACUUM_WALL_DB_PATH"] = old_path
if old_seed is not None:
os.environ["VACUUM_WALL_SEED_BUILTIN_ADMIN"] = old_seed
@pytest.fixture
def access_token(db_reset):
"""Mint a real access token for the seeded builtin admin."""
from lib.auth import generate_tokens
tokens = generate_tokens("admin", {"firewall": "rw"})
return tokens["access_token"]
class TestWsSnapshot:
def test_snapshot_sent_on_auth_connect(self, access_token):
"""A valid JWT subprotocol yields a full snapshot after auth."""
import asyncio
import daemon.server as server
store = State()
store.set("firewall", {"zones": {"public": {}}})
store.set("system", {"load": {"load1": 0.1}})
# Remaining subsystems stay None (not populated).
ws = AsyncMock()
ws.prepare = AsyncMock()
ws.send_json = AsyncMock()
request = MagicMock()
request.headers = {"Sec-WebSocket-Protocol": access_token}
with (
patch("aiohttp.web.WebSocketResponse", return_value=ws),
patch.object(server, "state_store", store),
):
asyncio.run(server._handle_ws(request))
ws.send_json.assert_awaited_once()
payload = ws.send_json.call_args[0][0]
assert payload["type"] == "snapshot"
data = payload["data"]
# Every subsystem key is present (push-stream: no `init`/`versions` shape).
for name in State.SUBSYSTEMS:
assert name in data
assert data["firewall"] == {"zones": {"public": {}}}
assert data["system"] == {"load": {"load1": 0.1}}
# Unpopulated subsystems are present but None (partial snapshot).
assert data["dnsmasq"] is None
assert data["wireguard"] is None
def test_no_snapshot_without_token(self):
"""Missing token -> 401 JSON, no WS is opened, no snapshot sent."""
import asyncio
import daemon.server as server
ws = AsyncMock()
request = MagicMock()
request.headers = {}
with patch("aiohttp.web.WebSocketResponse") as mock_ctor:
result = asyncio.run(server._handle_ws(request))
assert result.status == 401
mock_ctor.assert_not_called()
ws.send_json.assert_not_awaited()
def test_no_snapshot_on_invalid_token(self, access_token):
"""A token that fails validation -> 401, no snapshot sent."""
import asyncio
import daemon.server as server
ws = AsyncMock()
request = MagicMock()
request.headers = {"Sec-WebSocket-Protocol": access_token}
with (
patch("aiohttp.web.WebSocketResponse") as mock_ctor,
patch("lib.auth.validate_token", return_value=None),
):
result = asyncio.run(server._handle_ws(request))
assert result.status == 401
mock_ctor.assert_not_called()
ws.send_json.assert_not_awaited()