FastApi TestClient WebSockets | 如何測試正確使用線程和進程的 FastApi 應用程序？

Question

我想為我的 FastApi WebSocket 應用程序編寫測試，但是一個測試永遠運行並且不會停止，這會阻止下一個測試開始。

我在 FastApi 應用程序“啟動”期間調用了一個線程，該線程本身會產生進程。 我覺得使用線程和進程的組合確實在 TestClient 中產生了一些問題，因此測試永遠不會停止。

import uvicorn as uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import threading
import asyncio
from fastapi.testclient import TestClient
from multiprocessing import Process

app = FastAPI()

class ProcessHeavyCPU:
    def __init__(self): pass

    def run(self):  
        while True:
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within process")

class ThreadHandler(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        processes = []
        for _ in range(2):
            p = Process(target=ProcessHeavyCPU().run)
            processes.append(p)

        [x.start() for x in processes]

        while True:                     # <------ the test "test_example" runs forever
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within thread")



@app.on_event("startup")
async def startup_event():
    thread1 = ThreadHandler()
    thread1.daemon = True
    thread1.start()

async def handle_msg(websocket):
    await websocket.send_json({"msg": "New msg received"})

background_tasks = set()

@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            msg = await websocket.receive_text()
            task = asyncio.create_task(handle_msg(websocket))
            background_tasks.add(task)
            task.add_done_callback(background_tasks.discard)
    except WebSocketDisconnect:
        await websocket.close()


def test_example():         # <-- the test which does run forever
    with TestClient(app) as client:
        with client.websocket_connect("/") as websocket:
            websocket.send_json({
                "new_msg": "xyz"
            })
            resp = websocket.receive_json()
            assert resp["msg"] == "New msg received"
            print("finished test")


if __name__ == '__main__':
    uvicorn.run("main:app", host="0.0.0.0", port=8081, reload=True, access_log=False)

---

有什么想法可以解決這個問題，所以我可以在我的應用程序中使用線程和進程，同時使用舒適的 TestClient 進行 FastApi (Starlette) WebSocket 測試？

python -m pytest .\main.py

---

Python 3.10.5

文檔：

https://fastapi.tiangolo.com/advanced/testing-websockets/

Answer 1

我不確定這是否是您正在尋找的。 但是您的測試沒有停止的原因是因為您的線程沒有停止。 因此，僅在啟動時啟動線程，您可能還會在關閉時停止它。 所以它應該是這樣的：

import uvicorn as uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import threading
import asyncio
from fastapi.testclient import TestClient
from multiprocessing import Process

app = FastAPI()

class ProcessHeavyCPU:
    def __init__(self): pass

    def run(self):
        while True:
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within process")

class ThreadHandler(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.running = True

    def run(self):
        processes = []
        for _ in range(2):
            p = Process(target=ProcessHeavyCPU().run)
            processes.append(p)

        [x.start() for x in processes]

        while self.running:                     # <------ the test "test_example" runs forever
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within thread")

    def stop(self):
        self.running = False

thread1 = ThreadHandler()
@app.on_event("startup")
async def startup_event():
    thread1.daemon = True
    thread1.start()

@app.on_event("shutdown")
async def startup_event():
    thread1.stop()

async def handle_msg(websocket):
    await websocket.send_json({"msg": "New msg received"})

background_tasks = set()

@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            msg = await websocket.receive_text()
            task = asyncio.create_task(handle_msg(websocket))
            background_tasks.add(task)
            task.add_done_callback(background_tasks.discard)
    except WebSocketDisconnect:
        await websocket.close()


def test_example():         # <-- the test which does run forever
    with TestClient(app) as client:
        with client.websocket_connect("/") as websocket:
            websocket.send_json({
                "new_msg": "xyz"
            })
            resp = websocket.receive_json()
            assert resp["msg"] == "New msg received"
            print("finished test")


if __name__ == '__main__':
    uvicorn.run("main:app", host="0.0.0.0", port=8081, reload=True, access_log=False)