From cad7564c69d9c87420ce13742fe4fa59c018be7e Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Sat, 26 Aug 2023 14:58:47 +0800
Subject: [PATCH 1/5] Add WebSocket handling support for HTTP security
 dependencies

---
 fastapi/security/api_key.py             | 12 ++++++----
 fastapi/security/http.py                | 18 ++++++++-------
 fastapi/security/oauth2.py              | 13 ++++++-----
 fastapi/security/open_id_connect_url.py |  6 +++--
 fastapi/security/utils.py               | 29 ++++++++++++++++++++++++-
 5 files changed, 58 insertions(+), 20 deletions(-)

diff --git a/fastapi/security/api_key.py b/fastapi/security/api_key.py
index 8b2c5c080..83b404b80 100644
--- a/fastapi/security/api_key.py
+++ b/fastapi/security/api_key.py
@@ -2,8 +2,9 @@ from typing import Optional
 
 from fastapi.openapi.models import APIKey, APIKeyIn
 from fastapi.security.base import SecurityBase
+from fastapi.security.utils import handle_exc_for_ws
 from starlette.exceptions import HTTPException
-from starlette.requests import Request
+from starlette.requests import HTTPConnection
 from starlette.status import HTTP_403_FORBIDDEN
 
 
@@ -28,7 +29,8 @@ class APIKeyQuery(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.query_params.get(self.model.name)
         if not api_key:
             if self.auto_error:
@@ -57,7 +59,8 @@ class APIKeyHeader(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.headers.get(self.model.name)
         if not api_key:
             if self.auto_error:
@@ -86,7 +89,8 @@ class APIKeyCookie(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.cookies.get(self.model.name)
         if not api_key:
             if self.auto_error:
diff --git a/fastapi/security/http.py b/fastapi/security/http.py
index 8fc0aafd9..764dfac07 100644
--- a/fastapi/security/http.py
+++ b/fastapi/security/http.py
@@ -6,9 +6,9 @@ from fastapi.exceptions import HTTPException
 from fastapi.openapi.models import HTTPBase as HTTPBaseModel
 from fastapi.openapi.models import HTTPBearer as HTTPBearerModel
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import get_authorization_scheme_param
+from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
 from pydantic import BaseModel
-from starlette.requests import Request
+from starlette.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
 
@@ -35,8 +35,9 @@ class HTTPBase(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
+    @handle_exc_for_ws
     async def __call__(
-        self, request: Request
+        self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
         authorization = request.headers.get("Authorization")
         scheme, credentials = get_authorization_scheme_param(authorization)
@@ -64,9 +65,8 @@ class HTTPBasic(HTTPBase):
         self.realm = realm
         self.auto_error = auto_error
 
-    async def __call__(  # type: ignore
-        self, request: Request
-    ) -> Optional[HTTPBasicCredentials]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[HTTPBasicCredentials]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
         if self.realm:
@@ -110,8 +110,9 @@ class HTTPBearer(HTTPBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
+    @handle_exc_for_ws
     async def __call__(
-        self, request: Request
+        self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
         authorization = request.headers.get("Authorization")
         scheme, credentials = get_authorization_scheme_param(authorization)
@@ -145,8 +146,9 @@ class HTTPDigest(HTTPBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
+    @handle_exc_for_ws
     async def __call__(
-        self, request: Request
+        self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
         authorization = request.headers.get("Authorization")
         scheme, credentials = get_authorization_scheme_param(authorization)
diff --git a/fastapi/security/oauth2.py b/fastapi/security/oauth2.py
index e4c4357e7..6b4d20dcc 100644
--- a/fastapi/security/oauth2.py
+++ b/fastapi/security/oauth2.py
@@ -5,8 +5,8 @@ from fastapi.openapi.models import OAuth2 as OAuth2Model
 from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
 from fastapi.param_functions import Form
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import get_authorization_scheme_param
-from starlette.requests import Request
+from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
+from starlette.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
 # TODO: import from typing when deprecating Python 3.9
@@ -131,7 +131,8 @@ class OAuth2(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         if not authorization:
             if self.auto_error:
@@ -164,7 +165,8 @@ class OAuth2PasswordBearer(OAuth2):
             auto_error=auto_error,
         )
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
         if not authorization or scheme.lower() != "bearer":
@@ -210,7 +212,8 @@ class OAuth2AuthorizationCodeBearer(OAuth2):
             auto_error=auto_error,
         )
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
         if not authorization or scheme.lower() != "bearer":
diff --git a/fastapi/security/open_id_connect_url.py b/fastapi/security/open_id_connect_url.py
index 4e65f1f6c..390e99749 100644
--- a/fastapi/security/open_id_connect_url.py
+++ b/fastapi/security/open_id_connect_url.py
@@ -2,8 +2,9 @@ from typing import Optional
 
 from fastapi.openapi.models import OpenIdConnect as OpenIdConnectModel
 from fastapi.security.base import SecurityBase
+from fastapi.security.utils import handle_exc_for_ws
 from starlette.exceptions import HTTPException
-from starlette.requests import Request
+from starlette.requests import HTTPConnection
 from starlette.status import HTTP_403_FORBIDDEN
 
 
@@ -22,7 +23,8 @@ class OpenIdConnect(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[str]:
+    @handle_exc_for_ws
+    async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         if not authorization:
             if self.auto_error:
diff --git a/fastapi/security/utils.py b/fastapi/security/utils.py
index fa7a450b7..2a0849303 100644
--- a/fastapi/security/utils.py
+++ b/fastapi/security/utils.py
@@ -1,4 +1,10 @@
-from typing import Optional, Tuple
+from functools import wraps
+from typing import Any, Awaitable, Callable, Optional, Tuple, TypeVar
+
+from fastapi.exceptions import HTTPException, WebSocketException
+from starlette.requests import HTTPConnection
+from starlette.status import WS_1008_POLICY_VIOLATION
+from starlette.websockets import WebSocket
 
 
 def get_authorization_scheme_param(
@@ -8,3 +14,24 @@ def get_authorization_scheme_param(
         return "", ""
     scheme, _, param = authorization_header_value.partition(" ")
     return scheme, param
+
+
+_SecurityDepFunc = TypeVar(
+    "_SecurityDepFunc", bound=Callable[[Any, HTTPConnection], Awaitable]
+)
+
+
+def handle_exc_for_ws(func: _SecurityDepFunc) -> _SecurityDepFunc:
+    @wraps(func)
+    async def wrapper(self, request: HTTPConnection, *args, **kwargs):
+        try:
+            return await func(self, request, *args, **kwargs)
+        except HTTPException as e:
+            if not isinstance(request, WebSocket):
+                raise e
+            await request.accept()
+            raise WebSocketException(
+                code=WS_1008_POLICY_VIOLATION, reason=e.detail
+            ) from None
+
+    return wrapper  # type: ignore

From a10a186c889aac00fac581bbd5f79d8a8ada3d86 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Sat, 26 Aug 2023 15:17:30 +0800
Subject: [PATCH 2/5] Fix linting

---
 fastapi/security/http.py  | 4 +++-
 fastapi/security/utils.py | 6 +++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/fastapi/security/http.py b/fastapi/security/http.py
index 764dfac07..7c12f47ac 100644
--- a/fastapi/security/http.py
+++ b/fastapi/security/http.py
@@ -66,7 +66,9 @@ class HTTPBasic(HTTPBase):
         self.auto_error = auto_error
 
     @handle_exc_for_ws
-    async def __call__(self, request: HTTPConnection) -> Optional[HTTPBasicCredentials]:
+    async def __call__(  # type: ignore
+        self, request: HTTPConnection
+    ) -> Optional[HTTPBasicCredentials]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
         if self.realm:
diff --git a/fastapi/security/utils.py b/fastapi/security/utils.py
index 2a0849303..c40d526a7 100644
--- a/fastapi/security/utils.py
+++ b/fastapi/security/utils.py
@@ -17,15 +17,15 @@ def get_authorization_scheme_param(
 
 
 _SecurityDepFunc = TypeVar(
-    "_SecurityDepFunc", bound=Callable[[Any, HTTPConnection], Awaitable]
+    "_SecurityDepFunc", bound=Callable[[Any, HTTPConnection], Awaitable[Any]]
 )
 
 
 def handle_exc_for_ws(func: _SecurityDepFunc) -> _SecurityDepFunc:
     @wraps(func)
-    async def wrapper(self, request: HTTPConnection, *args, **kwargs):
+    async def wrapper(self: Any, request: HTTPConnection) -> Any:
         try:
-            return await func(self, request, *args, **kwargs)
+            return await func(self, request)
         except HTTPException as e:
             if not isinstance(request, WebSocket):
                 raise e

From 6a4ab615c85b49d91fc32de0d3789e7a1e6df88c Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Sat, 26 Aug 2023 16:20:34 +0800
Subject: [PATCH 3/5] Add tests for websocket with authorization

---
 fastapi/security/utils.py        |  3 ++-
 tests/test_security_http_base.py | 29 ++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/fastapi/security/utils.py b/fastapi/security/utils.py
index c40d526a7..3ddf56f1f 100644
--- a/fastapi/security/utils.py
+++ b/fastapi/security/utils.py
@@ -29,7 +29,8 @@ def handle_exc_for_ws(func: _SecurityDepFunc) -> _SecurityDepFunc:
         except HTTPException as e:
             if not isinstance(request, WebSocket):
                 raise e
-            await request.accept()
+            # close before accepted with result a HTTP 403 so the exception argument is ignored
+            # ref: https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
             raise WebSocketException(
                 code=WS_1008_POLICY_VIOLATION, reason=e.detail
             ) from None
diff --git a/tests/test_security_http_base.py b/tests/test_security_http_base.py
index 51928bafd..5c097939b 100644
--- a/tests/test_security_http_base.py
+++ b/tests/test_security_http_base.py
@@ -1,6 +1,8 @@
-from fastapi import FastAPI, Security
+import pytest
+from fastapi import FastAPI, Security, WebSocket
 from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBase
 from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
 
 app = FastAPI()
 
@@ -12,6 +14,16 @@ def read_current_user(credentials: HTTPAuthorizationCredentials = Security(secur
     return {"scheme": credentials.scheme, "credentials": credentials.credentials}
 
 
+@app.websocket("/users/timeline")
+async def read_user_timeline(
+    websocket: WebSocket, credentials: HTTPAuthorizationCredentials = Security(security)
+):
+    await websocket.accept()
+    await websocket.send_json(
+        {"scheme": credentials.scheme, "credentials": credentials.credentials}
+    )
+
+
 client = TestClient(app)
 
 
@@ -27,6 +39,21 @@ def test_security_http_base_no_credentials():
     assert response.json() == {"detail": "Not authenticated"}
 
 
+def test_security_http_base_with_ws():
+    with client.websocket_connect(
+        "/users/timeline", headers={"Authorization": "Other foobar"}
+    ) as websocket:
+        data = websocket.receive_json()
+        assert data == {"scheme": "Other", "credentials": "foobar"}
+
+
+def test_security_http_base_with_ws_no_credentials():
+    with pytest.raises(WebSocketDisconnect) as e:
+        with client.websocket_connect("/users/timeline"):
+            pass
+    assert e.value.reason == "Not authenticated"
+
+
 def test_openapi_schema():
     response = client.get("/openapi.json")
     assert response.status_code == 200, response.text

From 7edb2d9953ef1e0a6576085631ab4c147560bc8a Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Tue, 2 Apr 2024 23:01:26 +0800
Subject: [PATCH 4/5] Remove handle_exc_for_ws function from security utils

---
 fastapi/security/api_key.py             |  4 ----
 fastapi/security/http.py                |  6 +----
 fastapi/security/oauth2.py              |  5 +----
 fastapi/security/open_id_connect_url.py |  2 --
 fastapi/security/utils.py               | 30 +------------------------
 5 files changed, 3 insertions(+), 44 deletions(-)

diff --git a/fastapi/security/api_key.py b/fastapi/security/api_key.py
index 97f45555c..0456d2da8 100644
--- a/fastapi/security/api_key.py
+++ b/fastapi/security/api_key.py
@@ -2,7 +2,6 @@ from typing import Optional
 
 from fastapi.openapi.models import APIKey, APIKeyIn
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import handle_exc_for_ws
 from starlette.exceptions import HTTPException
 from starlette.requests import HTTPConnection
 from starlette.status import HTTP_403_FORBIDDEN
@@ -100,7 +99,6 @@ class APIKeyQuery(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.query_params.get(self.model.name)
         if not api_key:
@@ -196,7 +194,6 @@ class APIKeyHeader(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.headers.get(self.model.name)
         if not api_key:
@@ -292,7 +289,6 @@ class APIKeyCookie(APIKeyBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         api_key = request.cookies.get(self.model.name)
         if not api_key:
diff --git a/fastapi/security/http.py b/fastapi/security/http.py
index 8b8cc2632..1dbccbe8f 100644
--- a/fastapi/security/http.py
+++ b/fastapi/security/http.py
@@ -6,7 +6,7 @@ from fastapi.exceptions import HTTPException
 from fastapi.openapi.models import HTTPBase as HTTPBaseModel
 from fastapi.openapi.models import HTTPBearer as HTTPBearerModel
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
+from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import BaseModel
 from starlette.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
@@ -79,7 +79,6 @@ class HTTPBase(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(
         self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
@@ -185,7 +184,6 @@ class HTTPBasic(HTTPBase):
         self.realm = realm
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(  # type: ignore
         self, request: HTTPConnection
     ) -> Optional[HTTPBasicCredentials]:
@@ -300,7 +298,6 @@ class HTTPBearer(HTTPBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(
         self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
@@ -403,7 +400,6 @@ class HTTPDigest(HTTPBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(
         self, request: HTTPConnection
     ) -> Optional[HTTPAuthorizationCredentials]:
diff --git a/fastapi/security/oauth2.py b/fastapi/security/oauth2.py
index 35c22b5bb..e27edb10f 100644
--- a/fastapi/security/oauth2.py
+++ b/fastapi/security/oauth2.py
@@ -5,7 +5,7 @@ from fastapi.openapi.models import OAuth2 as OAuth2Model
 from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
 from fastapi.param_functions import Form
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
+from fastapi.security.utils import get_authorization_scheme_param
 from starlette.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
@@ -376,7 +376,6 @@ class OAuth2(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         if not authorization:
@@ -471,7 +470,6 @@ class OAuth2PasswordBearer(OAuth2):
             auto_error=auto_error,
         )
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
@@ -582,7 +580,6 @@ class OAuth2AuthorizationCodeBearer(OAuth2):
             auto_error=auto_error,
         )
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
diff --git a/fastapi/security/open_id_connect_url.py b/fastapi/security/open_id_connect_url.py
index 9dbd8d8ea..0e4643ef5 100644
--- a/fastapi/security/open_id_connect_url.py
+++ b/fastapi/security/open_id_connect_url.py
@@ -2,7 +2,6 @@ from typing import Optional
 
 from fastapi.openapi.models import OpenIdConnect as OpenIdConnectModel
 from fastapi.security.base import SecurityBase
-from fastapi.security.utils import handle_exc_for_ws
 from starlette.exceptions import HTTPException
 from starlette.requests import HTTPConnection
 from starlette.status import HTTP_403_FORBIDDEN
@@ -73,7 +72,6 @@ class OpenIdConnect(SecurityBase):
         self.scheme_name = scheme_name or self.__class__.__name__
         self.auto_error = auto_error
 
-    @handle_exc_for_ws
     async def __call__(self, request: HTTPConnection) -> Optional[str]:
         authorization = request.headers.get("Authorization")
         if not authorization:
diff --git a/fastapi/security/utils.py b/fastapi/security/utils.py
index 3ddf56f1f..fa7a450b7 100644
--- a/fastapi/security/utils.py
+++ b/fastapi/security/utils.py
@@ -1,10 +1,4 @@
-from functools import wraps
-from typing import Any, Awaitable, Callable, Optional, Tuple, TypeVar
-
-from fastapi.exceptions import HTTPException, WebSocketException
-from starlette.requests import HTTPConnection
-from starlette.status import WS_1008_POLICY_VIOLATION
-from starlette.websockets import WebSocket
+from typing import Optional, Tuple
 
 
 def get_authorization_scheme_param(
@@ -14,25 +8,3 @@ def get_authorization_scheme_param(
         return "", ""
     scheme, _, param = authorization_header_value.partition(" ")
     return scheme, param
-
-
-_SecurityDepFunc = TypeVar(
-    "_SecurityDepFunc", bound=Callable[[Any, HTTPConnection], Awaitable[Any]]
-)
-
-
-def handle_exc_for_ws(func: _SecurityDepFunc) -> _SecurityDepFunc:
-    @wraps(func)
-    async def wrapper(self: Any, request: HTTPConnection) -> Any:
-        try:
-            return await func(self, request)
-        except HTTPException as e:
-            if not isinstance(request, WebSocket):
-                raise e
-            # close before accepted with result a HTTP 403 so the exception argument is ignored
-            # ref: https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
-            raise WebSocketException(
-                code=WS_1008_POLICY_VIOLATION, reason=e.detail
-            ) from None
-
-    return wrapper  # type: ignore

From 6846708f007b615f99453dc2730596eebfe08fd8 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Wed, 3 Apr 2024 13:58:24 +0800
Subject: [PATCH 5/5] Refactor test script with `auto_error=False` case.

---
 tests/test_security_http_base.py          | 29 +----------------------
 tests/test_security_http_base_optional.py | 29 ++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/tests/test_security_http_base.py b/tests/test_security_http_base.py
index 5c097939b..51928bafd 100644
--- a/tests/test_security_http_base.py
+++ b/tests/test_security_http_base.py
@@ -1,8 +1,6 @@
-import pytest
-from fastapi import FastAPI, Security, WebSocket
+from fastapi import FastAPI, Security
 from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBase
 from fastapi.testclient import TestClient
-from starlette.websockets import WebSocketDisconnect
 
 app = FastAPI()
 
@@ -14,16 +12,6 @@ def read_current_user(credentials: HTTPAuthorizationCredentials = Security(secur
     return {"scheme": credentials.scheme, "credentials": credentials.credentials}
 
 
-@app.websocket("/users/timeline")
-async def read_user_timeline(
-    websocket: WebSocket, credentials: HTTPAuthorizationCredentials = Security(security)
-):
-    await websocket.accept()
-    await websocket.send_json(
-        {"scheme": credentials.scheme, "credentials": credentials.credentials}
-    )
-
-
 client = TestClient(app)
 
 
@@ -39,21 +27,6 @@ def test_security_http_base_no_credentials():
     assert response.json() == {"detail": "Not authenticated"}
 
 
-def test_security_http_base_with_ws():
-    with client.websocket_connect(
-        "/users/timeline", headers={"Authorization": "Other foobar"}
-    ) as websocket:
-        data = websocket.receive_json()
-        assert data == {"scheme": "Other", "credentials": "foobar"}
-
-
-def test_security_http_base_with_ws_no_credentials():
-    with pytest.raises(WebSocketDisconnect) as e:
-        with client.websocket_connect("/users/timeline"):
-            pass
-    assert e.value.reason == "Not authenticated"
-
-
 def test_openapi_schema():
     response = client.get("/openapi.json")
     assert response.status_code == 200, response.text
diff --git a/tests/test_security_http_base_optional.py b/tests/test_security_http_base_optional.py
index dd4d76843..4da9e82bc 100644
--- a/tests/test_security_http_base_optional.py
+++ b/tests/test_security_http_base_optional.py
@@ -1,6 +1,6 @@
 from typing import Optional
 
-from fastapi import FastAPI, Security
+from fastapi import FastAPI, Security, WebSocket
 from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBase
 from fastapi.testclient import TestClient
 
@@ -18,6 +18,19 @@ def read_current_user(
     return {"scheme": credentials.scheme, "credentials": credentials.credentials}
 
 
+@app.websocket("/users/timeline")
+async def read_user_timeline(
+    websocket: WebSocket,
+    credentials: Optional[HTTPAuthorizationCredentials] = Security(security),
+):
+    await websocket.accept()
+    await websocket.send_json(
+        {"scheme": credentials.scheme, "credentials": credentials.credentials}
+        if credentials
+        else {"msg": "Create an account first"}
+    )
+
+
 client = TestClient(app)
 
 
@@ -33,6 +46,20 @@ def test_security_http_base_no_credentials():
     assert response.json() == {"msg": "Create an account first"}
 
 
+def test_security_http_base_with_ws():
+    with client.websocket_connect(
+        "/users/timeline", headers={"Authorization": "Other foobar"}
+    ) as websocket:
+        data = websocket.receive_json()
+        assert data == {"scheme": "Other", "credentials": "foobar"}
+
+
+def test_security_http_base_with_ws_no_credentials():
+    with client.websocket_connect("/users/timeline") as websocket:
+        data = websocket.receive_json()
+        assert data == {"msg": "Create an account first"}
+
+
 def test_openapi_schema():
     response = client.get("/openapi.json")
     assert response.status_code == 200, response.text