其實代碼沒有問題，但是我們忽略了一點，就是我們在正常的開發中，肯定是遇到這樣的情況，我們頻繁的有客戶端鏈接，斷開鏈接，我們需要統一的管理起來，那么我們應該如何去管理呢，其實這個時候，我們要去宣告一個類去管理我們的這些鏈接，我們應該如何優化呢，
定義一個鏈接管理類，處理我們所有的鏈接，
class ConnectionManager: def __init__(self): # 存放**的鏈接 self.active_connections: List[Dict[str, WebSocket]] = [] async def connect(self, user: str, ws: WebSocket): # 鏈接 await ws.accept() self.active_connections.append({"user": user, "ws": ws}) def disconnect(self, user: str, ws: WebSocket): # 關閉時 移除ws物件 self.active_connections.remove({"user": user, "ws": ws})
我們增加了鏈接，移除鏈接的操作，那么我們垓下我們的工程代碼
manager = ConnectionManager() @app.websocket("/items/ws") async def websocket_endpoint( websocket: WebSocket, cookie_or_token: str = Depends(get_cookie_or_token), ): await manager.connect(cookie_or_token, websocket) try: while True: data = await websocket.receive_text() await websocket.send_text(f"訊息是: {data}") except WebSocketDisconnect as e: manager.disconnect(cookie_or_token, websocket)
這樣我們在連接處理的時候就可以正常處理了，之前報錯是因為我們沒有正常的關閉連接導致的，那么我們這個時候再來去看下測驗代碼
def test_websocket(): client = TestClient(app) with client.websocket_connect("/items/ws?token="+"leizishuoceshi") as websocket: websocket.send_text("Hello WebSocket") data = websocket.receive_text() print(data) assert str(data) =="訊息是: Hello WebSocket"
此時我們再去執行，發現代碼應該不會報錯
我們在用例優化下
class FastApiTestWeb(unittest.TestCase): def setUp(self) -> None: self.client = TestClient(app) def tearDown(self) -> None: self.client = None def test_websocket(self): with self.client.websocket_connect("/items/ws?token=" + "leizishuoceshi") as websocket: websocket.send_text("Hello WebSocket") data = websocket.receive_text() print(data) assert str(data) == "訊息是: Hello WebSocket" def test_websocket_two(self): with self.client.websocket_connect("/items/ws?token=" + "leizishuoceshi") as websocket: websocket.send_text("Hello 123") data = websocket.receive_text() print(data) assert str(data) == "訊息是: Hello 123" if __name__ == "__main__": unittest.main()
這樣我們的一個測驗用例就更加的完整了，我們執行正常是沒有報錯的
我們想要看下代碼的覆寫率，應該如何看呢，我是用的coverage，
然后我們再去report，
我們想看html測驗報告，可以運行下 coverage html，
然后看index.html
因為我的main.py還有其他的方法，我們還需要點進去看我們對應方法的覆寫率，
我們可以看到，我們的關鍵代碼還差了一行沒有覆寫到，就是不帶token的訪問，我們在增加下一條case去覆寫下，
def test_websocket_notoken(self): try: with self.client.websocket_connect("/items/ws") as websocket: websocket.send_text("Hello 123") data = websocket.receive_text() self.assertEqual("訊息是: Hello 123",data) except: self.assertTrue(True)
覆寫了這行代碼，但是我們用例需要判斷這個例外，
到這里，我們對于WebSockets介面測驗完畢，但是還有問題，我們真正的聊天中，還需要上線進行通知，下線進行通知，我們應該如何實作呢，且聽下回分解，
文章首發在公眾號，歡迎關注，
轉載請註明出處，本文鏈接：https://www.uj5u.com/houduan/335068.html
標籤：Python