test(server): cover WebSocket authentication and authorization

The WebSocket scope is now authenticated and authorized like HTTP, so update
the route-authorization test that asserted WebSocket passthrough to instead
verify that a permitted route connects and a non-permitted one is rejected with
a close frame. Add AuthenticationMiddleware tests covering WebSocket handshakes
with a missing token, an invalid token, and a valid token.

Signed-off-by: Charlie Doern <cdoern@redhat.com>

Author: cdoern

tests/unit/server/test_auth.py

@@ -10,7 +10,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.testclient import TestClient
 
 from ogx.core.datatypes import (
@@ -96,6 +96,32 @@ def http_client(http_app):
     return TestClient(http_app)
 
 
+@pytest.fixture
+def ws_app(mock_auth_endpoint):
+    app = FastAPI()
+    auth_config = AuthenticationConfig(
+        provider_config=CustomAuthConfig(
+            type=AuthProviderType.CUSTOM,
+            endpoint=mock_auth_endpoint,
+        ),
+        access_policy=[],
+    )
+    app.add_middleware(AuthenticationMiddleware, auth_config=auth_config, impls={})
+
+    @app.websocket("/ws")
+    async def ws_endpoint(websocket: WebSocket):
+        await websocket.accept()
+        await websocket.send_text("authenticated")
+        await websocket.close()
+
+    return app
+
+
+@pytest.fixture
+def ws_client(ws_app):
+    return TestClient(ws_app)
+
+
 @pytest.fixture
 def mock_scope():
     return {
@@ -226,6 +252,26 @@ def test_http_auth_service_error(http_client, valid_api_key, suppress_auth_error
     assert "Authentication service error" in response.json()["error"]["message"]
 
 
+# WebSocket Endpoint Tests
+def test_websocket_missing_auth_header_rejected(ws_client, suppress_auth_errors):
+    with pytest.raises(WebSocketDisconnect):
+        with ws_client.websocket_connect("/ws") as websocket:
+            websocket.receive_text()
+
+
+@patch("httpx.AsyncClient.post", new=mock_post_failure)
+def test_websocket_invalid_authentication_rejected(ws_client, invalid_api_key, suppress_auth_errors):
+    with pytest.raises(WebSocketDisconnect):
+        with ws_client.websocket_connect("/ws", headers={"Authorization": f"Bearer {invalid_api_key}"}) as websocket:
+            websocket.receive_text()
+
+
+@patch("httpx.AsyncClient.post", new=mock_post_success)
+def test_websocket_valid_authentication_accepted(ws_client, valid_api_key):
+    with ws_client.websocket_connect("/ws", headers={"Authorization": f"Bearer {valid_api_key}"}) as websocket:
+        assert websocket.receive_text() == "authenticated"
+
+
 def test_http_auth_request_payload(http_client, valid_api_key, mock_auth_endpoint, suppress_auth_errors):
     with patch("httpx.AsyncClient.post") as mock_post:
         mock_response = MockResponse(200, {"message": "Authentication successful"})

tests/unit/server/test_route_auth.py

@@ -343,42 +343,44 @@ async def test_rule_order_matters(developer_user):
     assert not middleware._is_route_allowed("/v1/models/list", developer_user)
 
 
-async def test_websocket_passthrough():
-    """Test that websocket requests pass through without blocking"""
+async def test_websocket_route_authorization():
+    """WebSocket handshakes are subject to the same route policy as HTTP."""
     route_policy = [
         RouteAccessRule(
-            permit=RouteScope(paths="/v1/chat/completions"),
-            when="user with developer in roles",
+            permit=RouteScope(paths="/v1/responses"),
+            description="Allow the Responses WebSocket route",
         )
+        # All other routes denied by default (no matching rule)
     ]
 
-    app = FastAPI()
-    middleware = RouteAuthorizationMiddleware(app, route_policy)
-
-    # Mock websocket scope
-    scope = {
-        "type": "websocket",
-        "path": "/ws",
-    }
-
-    # Track if next middleware was called
-    called = False
-
-    async def mock_app(scope, receive, send):
-        nonlocal called
-        called = True
-
     async def receive():
         return {}
 
-    async def send(msg):
-        pass
+    async def run(path: str) -> tuple[bool, list[dict]]:
+        called = False
+        sent: list[dict] = []
 
-    middleware.app = mock_app
-    await middleware(scope, receive, send)
+        async def mock_app(scope, receive, send):
+            nonlocal called
+            called = True
+
+        async def send(msg):
+            sent.append(msg)
 
-    # Websocket requests should pass through
+        middleware = RouteAuthorizationMiddleware(mock_app, route_policy)
+        await middleware({"type": "websocket", "path": path}, receive, send)
+        return called, sent
+
+    # Permitted route reaches the app.
+    called, sent = await run("/v1/responses")
     assert called
+    assert sent == []
+
+    # A route with no matching permit rule is rejected with a close frame
+    # instead of being allowed through unauthenticated.
+    called, sent = await run("/v1/admin/reset")
+    assert not called
+    assert sent == [{"type": "websocket.close", "code": 4403}]
 
 
 async def test_route_blocking_without_auth():