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] 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