Search code examples
pythonpytestfastapiuvicorn

How to terminate a Uvicorn + FastAPI application cleanly with workers >= 2 when testing with pytest


I have an application written with Uvicorn + FastAPI. I am testing the response time using PyTest.

Referring to How to start a Uvicorn + FastAPI in background when testing with PyTest, I wrote the test. However, I found the application process alive after completing the test when workers >= 2.

I want to terminate the application process cleanly at the end of the test.

Do you have any idea?

The details are as follows.

Environment

Libraries

  • fastapi == 0.68.0
  • uvicorn == 0.14.0
  • requests == 2.26.0
  • pytest == 6.2.4

Sample Codes

  • Application: main.py
    from fastapi import FastAPI
    
    app = FastAPI()
    
    @app.get("/")
    def hello_world():
        return "hello world"
    
  • Test: test_main.py
    from multiprocessing import Process
    import pytest
    import requests
    import time
    import uvicorn
    
    HOST = "127.0.0.1"
    PORT = 8765
    WORKERS = 1
    
    
    def run_server(host: str, port: int, workers: int, wait: int = 15) -> Process:
        proc = Process(
            target=uvicorn.run,
            args=("main:app",),
            kwargs={
                "host": host,
                "port": port,
                "workers": workers,
            },
        )
        proc.start()
        time.sleep(wait)
        assert proc.is_alive()
        return proc
    
    
    def shutdown_server(proc: Process):
        proc.terminate()
        for _ in range(5):
            if proc.is_alive():
                time.sleep(5)
            else:
                return
        else:
            raise Exception("Process still alive")
    
    
    def check_response(host: str, port: int):
        assert requests.get(f"http://{host}:{port}").text == '"hello world"'
    
    
    def check_response_time(host: str, port: int, tol: float = 1e-2):
        s = time.time()
        requests.get(f"http://{host}:{port}")
        e = time.time()
        assert e-s < tol
    
    
    @pytest.fixture(scope="session")
    def server():
        proc = run_server(HOST, PORT, WORKERS)
        try:
            yield
        finally:
            shutdown_server(proc)
    
    
    def test_main(server):
        check_response(HOST, PORT)
        check_response_time(HOST, PORT)
        check_response(HOST, PORT)
        check_response_time(HOST, PORT)
    

Execution Result

$ curl http://localhost:8765
curl: (7) Failed to connect to localhost port 8765: Connection refused
$ pytest test_main.py
=============== test session starts =============== platform win32 -- Python 3.7.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1
rootdir: .\
collected 1 item

test_main.py .                                                                                                                                                                                                                         [100%]

=============== 1 passed in 20.23s ===============
$ curl http://localhost:8765
curl: (7) Failed to connect to localhost port 8765: Connection refused
$ sed -i -e "s/WORKERS = 1/WORKERS = 3/g" test_main.py
$ curl http://localhost:8765
curl: (7) Failed to connect to localhost port 8765: Connection refused
$ pytest test_main.py
=============== test session starts =============== platform win32 -- Python 3.7.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1
rootdir: .\
collected 1 item

test_main.py .                                                                                                                                                                                                                         [100%]

=============== 1 passed in 20.21s ===============
$ curl http://localhost:8765
"hello world"

$ # Why is localhost:8765 still alive?

Solution

  • I have found a solution myself.

    Thanks > https://stackoverflow.com/a/27034438/16567832

    Solution

    After install psutil by pip install psutil, update test_main.py

    from multiprocessing import Process
    import psutil
    import pytest
    import requests
    import time
    import uvicorn
    
    HOST = "127.0.0.1"
    PORT = 8765
    WORKERS = 3
    
    
    def run_server(host: str, port: int, workers: int, wait: int = 15) -> Process:
        proc = Process(
            target=uvicorn.run,
            args=("main:app",),
            kwargs={
                "host": host,
                "port": port,
                "workers": workers,
            },
        )
        proc.start()
        time.sleep(wait)
        assert proc.is_alive()
        return proc
    
    
    def shutdown_server(proc: Process):
    
        ##### SOLUTION #####
        pid = proc.pid
        parent = psutil.Process(pid)
        for child in parent.children(recursive=True):
            child.kill()
        ##### SOLUTION END ####
    
        proc.terminate()
        for _ in range(5):
            if proc.is_alive():
                time.sleep(5)
            else:
                return
        else:
            raise Exception("Process still alive")
    
    
    def check_response(host: str, port: int):
        assert requests.get(f"http://{host}:{port}").text == '"hello world"'
    
    
    def check_response_time(host: str, port: int, tol: float = 1e-2):
        s = time.time()
        requests.get(f"http://{host}:{port}")
        e = time.time()
        assert e-s < tol
    
    
    @pytest.fixture(scope="session")
    def server():
        proc = run_server(HOST, PORT, WORKERS)
        try:
            yield
        finally:
            shutdown_server(proc)
    
    
    def test_main(server):
        check_response(HOST, PORT)
        check_response_time(HOST, PORT)
        check_response(HOST, PORT)
        check_response_time(HOST, PORT)
    

    Execution Result

    $ curl http://localhost:8765
    curl: (7) Failed to connect to localhost port 8765: Connection refused
    $ pytest test_main.py
    ================== test session starts ================== platform win32 -- Python 3.7.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1
    rootdir: .\
    collected 1 item
    
    test_main.py .                                                                                                                                                                                                                         [100%]
    
    ================== 1 passed in 20.24s ==================
    $ curl http://localhost:8765
    curl: (7) Failed to connect to localhost port 8765: Connection refused