Fix ws conn handler, add ws tests

- Fixes permission issue in connection handler event loop
- Adds tests for the websockets router and conn handler

Author: mawaeg

api/utils/websocket_connection_handler.py

@@ -3,19 +3,19 @@
 from fastapi import HTTPException, WebSocket
 from sqlalchemy.ext.asyncio import AsyncEngine
 
-from api.models.database_models import DatabaseModelBase, DBUser
+from api.models.database_models import DBUser, SensorData, SensorState
 from api.utils.database import get_engine, get_session
 from api.utils.permissions import get_user_read_permissions
 
 
 class WebsocketHandler:
     def __init__(self):
         self._connections: dict[int, tuple[DBUser, WebSocket]] = {}
-        self._message_queue: asyncio.Queue[DatabaseModelBase] = asyncio.Queue()
+        self._message_queue: asyncio.Queue[SensorState | SensorData] = asyncio.Queue()
 
     def add(self, user: DBUser, websocket: WebSocket) -> bool:
         if not user.id:
-            return False
+            return False  # pragma: no cover: Just for type safety, can not really happen (Should we remove the check?)
         if self._connections.get(user.id):
             return False
 
@@ -24,10 +24,10 @@ def add(self, user: DBUser, websocket: WebSocket) -> bool:
 
     def remove(self, user: DBUser) -> bool:
         if not user.id:
-            return False
+            return False  # pragma: no cover: Just for type safety, can not really happen (Should we remove the check?)
         if self._connections.pop(user.id, None):
             return True
-        return False
+        return False  # pragma: no cover: Just for type safety, shouldn't happen
 
     async def event_loop(self):
         engine: AsyncEngine = await get_engine()
@@ -36,14 +36,13 @@ async def event_loop(self):
             async for session in get_session(engine):
                 for user, websocket in self._connections.values():
                     try:
-                        if event.id:
-                            await get_user_read_permissions(session, user, event.id)
-                            await websocket.send_text(event.model_dump_json())
-                    except HTTPException:
+                        await get_user_read_permissions(session, user, event.sensor_id)
+                        await websocket.send_text(event.model_dump_json())
+                    except HTTPException:  # pragma: no cover: For now no tests, permission handling is tested elsewhere
                         pass
                 self._message_queue.task_done()
 
-    async def add_event(self, data: DatabaseModelBase):
+    async def add_event(self, data: SensorState | SensorData):
         await self._message_queue.put(data)
 
 

tests/router_tests/test_websocket.py

@@ -0,0 +1,86 @@
+import asyncio
+from typing import Any
+
+import pytest
+from fastapi import WebSocketDisconnect
+from fastapi.testclient import TestClient
+
+from api.main import app
+from api.models.database_models import Sensor, SensorData
+from api.utils.database import get_engine
+from api.utils.websocket_connection_handler import get_websocket_handler
+from tests.utils.fake_db import override_get_engine
+from tests.utils.fixtures import superuser_token, token
+from tests.utils.sensor_utils import create_sensor
+
+app.dependency_overrides[get_engine] = override_get_engine
+client: TestClient = TestClient(app)
+
+
+def test_ws_connect_no_authentication():
+    try:
+        with client.websocket_connect("/ws") as _:
+            return
+    except WebSocketDisconnect as exception:
+        assert exception.code == 1008
+
+
+def test_ws_connect_no_bearer():
+    try:
+        with client.websocket_connect("/ws", headers={"Authorization": "1234"}) as _:
+            return
+    except WebSocketDisconnect as exception:
+        assert exception.code == 1008
+
+
+def test_ws_connect_wrong_token(token: str):
+    try:
+        with client.websocket_connect("/ws", headers={"Authorization": f"Bearer {token}1234"}) as _:
+            return
+    except WebSocketDisconnect as exception:
+        assert exception.code == 1008
+
+
+@pytest.mark.asyncio
+async def test_ws_connect(superuser_token: str):
+    ws_handler = get_websocket_handler()
+    task = asyncio.get_event_loop().create_task(get_websocket_handler().event_loop())
+
+    sensor: Sensor = await create_sensor(name="TestSensorWS")
+
+    with client.websocket_connect("/ws", headers={"Authorization": f"Bearer {superuser_token}"}) as websocket:
+        data = websocket.receive_json()
+        assert data == {"message": "Hello World!"}
+
+        data: dict[str, Any] = {
+            "id": -1,
+            "temperature": 32.1,
+            "humidity": 56.78,
+            "pressure": 123.45,
+            "voltage": 3.45,
+            "sensor_id": sensor.id,
+        }
+        sensor_data: SensorData = SensorData(**data)
+        await ws_handler.add_event(sensor_data)
+
+        # As the websocket.receive_json() is synchronous and TestClient does not really work with async ws,
+        # the delay is need, so that the ws_handler can run its loop.
+        await asyncio.sleep(0.1)
+
+        received_data = websocket.receive_json()
+        assert data.get("id") == received_data.get("id")
+        assert data.get("temperature") == received_data.get("temperature")
+    task.cancel()
+
+
+# ToDo Investigate performance issue
+def test_ws_connect_double_connection(token: str):
+    with client.websocket_connect("/ws", headers={"Authorization": f"Bearer {token}"}) as websocket:
+        data = websocket.receive_json()
+        assert data == {"message": "Hello World!"}
+
+        try:
+            with client.websocket_connect("/ws", headers={"Authorization": f"Bearer {token}"}) as _:
+                pass
+        except WebSocketDisconnect as exception:
+            assert exception.code == 1008

tests/utils_tests/test_security.py

@@ -1,18 +1,19 @@
 import datetime
 
 import pytest
-from fastapi import HTTPException
+from fastapi import HTTPException, WebSocketException
+from pytest_mock import MockerFixture
 from sqlmodel import delete
 
 from api.models.database_models import DBUser
 from api.utils.http_exceptions import INVALID_CREDENTIALS
-from api.utils.security import create_access_token, get_current_user
+from api.utils.security import create_access_token, get_current_user, get_current_user_ws
 from tests.utils.fake_db import async_fake_session_maker, initialize_fake_database
 
 
 @pytest.mark.parametrize("token_data", [{}, {"sub": "fake_username"}])
 @pytest.mark.asyncio
-async def test_get_current_user_invalid_token(token_data: str):
+async def test_get_current_user_invalid_token(token_data: dict):
     await initialize_fake_database()
     async with async_fake_session_maker() as session:
         await session.execute(delete(DBUser))
@@ -24,3 +25,23 @@ async def test_get_current_user_invalid_token(token_data: str):
             await get_current_user(access_token, session)
 
         assert exception.value == INVALID_CREDENTIALS
+
+
+@pytest.mark.parametrize("token_data", [{}, {"sub": "fake_username"}])
+@pytest.mark.asyncio
+async def test_get_current_user_ws_invalid_token(token_data: dict, mocker: MockerFixture):
+    await initialize_fake_database()
+    async with async_fake_session_maker() as session:
+        await session.execute(delete(DBUser))
+        await session.commit()
+
+        access_token = create_access_token(data=token_data, expires_delta=datetime.timedelta(minutes=1))
+
+        jwt_decode_mock = mocker.patch("api.utils.security.jwt.decode")
+        jwt_decode_mock.return_value = token_data
+
+        with pytest.raises(WebSocketException) as exception:
+            await get_current_user_ws(session, f"Bearer {access_token}")
+
+        # assert jwt_decode_mock.assert_called_once()
+        assert exception.value.code == 1008