使用 pytest 进行测试时，如何使用工人 >= 2 干净地终止 Uvicorn + FastAPI 应用程序

Question

我有一个用 Uvicorn + FastAPI 编写的应用程序。 我正在使用 PyTest 测试响应时间。

参考How to start a Uvicorn + FastAPI in background when testing with PyTest，我写了测试。 但是，当worker >= 2时，我完成测试后发现应用程序还活着。

我想在测试结束时干净地终止应用程序进程。

你有什么想法吗？

详情如下。

环境

• Windows 10
• Bash 4.4.23 (https://cmder.net/)
• python 3.7.5

库

• fastapi == 0.68.0
• uvicorn == 0.14.0
• 请求 == 2.26.0
• pytest == 6.2.4

示例代码

• 应用程序：main.py

  from fastapi import FastAPI
  
  app = FastAPI()
  
  @app.get("/")
  def hello_world():
      return "hello world"
  

• 测试：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)
  

执行结果

$ 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?

Answer 1

我自己找到了解决办法。

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

解决方案

pip install psutil安装psutil后，更新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)

执行结果

$ 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