fix(api): defer socket.accept until after plugin resolution, parametrize tests

Aureliolo · claude · Aureliolo · commit 557764409fe0 · 2026-03-18T22:14:34.000+01:00
- Move socket.accept() after ChannelsPlugin resolution and subscribe()
  so infrastructure failures never accept-then-immediately-close
- Merge test_ws_rejects_missing_ticket and test_ws_rejects_invalid_ticket
  into a single parametrized test_ws_rejects_bad_ticket

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/synthorg/api/controllers/ws.py b/src/synthorg/api/controllers/ws.py
@@ -215,17 +215,6 @@ async def ws_handler(
     if not await _check_ws_role(socket, user):
         return
 
-    socket.scope["user"] = user
-    await socket.accept()
-    logger.info(
-        API_WS_CONNECTED,
-        client=str(socket.client),
-        user_id=user.user_id,
-    )
-
-    subscribed: set[str] = set()
-    filters: dict[str, dict[str, str]] = {}
-
     # Resolve ChannelsPlugin by iterating app.plugins -- the same
     # pattern used by get_channels_plugin() in channels.py.
     # Litestar's DI does not reliably inject plugin instances into
@@ -244,8 +233,23 @@ async def ws_handler(
         )
         await socket.close(code=1011, reason="Internal error")
         return
+
     subscriber = await channels_plugin.subscribe(list(ALL_CHANNELS))
 
+    # Accept only after auth, role check, plugin resolution, and
+    # subscription all succeed -- avoid accepting then immediately
+    # closing on infrastructure failures.
+    socket.scope["user"] = user
+    await socket.accept()
+    logger.info(
+        API_WS_CONNECTED,
+        client=str(socket.client),
+        user_id=user.user_id,
+    )
+
+    subscribed: set[str] = set()
+    filters: dict[str, dict[str, str]] = {}
+
     async def _event_callback(event_data: bytes) -> None:
         await _on_event(event_data, subscribed, filters, socket)
 
diff --git a/tests/unit/api/controllers/test_ws.py b/tests/unit/api/controllers/test_ws.py
@@ -281,11 +281,20 @@ def test_ws_close_codes_in_application_range(self) -> None:
         assert 4000 <= _WS_CLOSE_AUTH_FAILED <= 4999
         assert 4000 <= _WS_CLOSE_FORBIDDEN <= 4999
 
-    def test_ws_rejects_missing_ticket(
+    @pytest.mark.parametrize(
+        ("url", "scenario"),
+        [
+            ("/api/v1/ws", "missing_ticket"),
+            ("/api/v1/ws?ticket=bogus-ticket", "invalid_ticket"),
+        ],
+    )
+    def test_ws_rejects_bad_ticket(
         self,
         test_client: TestClient[Any],
+        url: str,
+        scenario: str,
     ) -> None:
-        """WS connection without ?ticket= is rejected with code 4001.
+        """WS connection with missing or invalid ticket is rejected.
 
         Verifying the close code (not just WebSocketDisconnect) ensures
         the rejection comes from the handler's ticket validation -- not
@@ -296,24 +305,13 @@ def test_ws_rejects_missing_ticket(
 
         with (
             pytest.raises(WebSocketDisconnect) as exc_info,
-            test_client.websocket_connect("/api/v1/ws"),
+            test_client.websocket_connect(url),
         ):
             pass
-        assert exc_info.value.code == _WS_CLOSE_AUTH_FAILED
-
-    def test_ws_rejects_invalid_ticket(
-        self,
-        test_client: TestClient[Any],
-    ) -> None:
-        """WS connection with a bogus ticket is rejected with code 4001."""
-        from litestar.exceptions import WebSocketDisconnect
-
-        with (
-            pytest.raises(WebSocketDisconnect) as exc_info,
-            test_client.websocket_connect("/api/v1/ws?ticket=bogus-ticket"),
-        ):
-            pass
-        assert exc_info.value.code == _WS_CLOSE_AUTH_FAILED
+        assert exc_info.value.code == _WS_CLOSE_AUTH_FAILED, (
+            f"Expected close code {_WS_CLOSE_AUTH_FAILED} for "
+            f"{scenario}, got {exc_info.value.code}"
+        )
 
     def test_ws_accepts_valid_ticket(
         self,
