FastAPI 學習之路（四十九）WebSockets（五）修復接口測試中的問題

其實代碼沒有問題，但是我們忽略了一點，就是我們在正常的開發中，肯定是遇到這樣的情況，我們頻繁的有客戶端鏈接，斷開鏈接，我們需要統一的管理起來，那麼我們應該如何去管理呢，其實這個時候，我們要去聲明一個類去管理我們的這些鏈接。我們應該如何優化呢。

      定義一個鏈接管理類，處理我們所有的鏈接。

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接口測試完畢，但是還有問題，我們真正的聊天中，還需要上線進行通知，下線進行通知，我們應該如何實現呢，且聽下回分解。

文章首發在公衆號，歡迎關注。