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

阿新 • • 發佈：2021-10-24

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

定義一個連結管理類，處理我們所有的連結。

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)