简体   繁体   中英

How to test a FastAPI api endpoint that consumes images?

I am using pytest to test a FastAPI endpoint that gets in input an image in binary format as in

@app.post("/analyse")
async def analyse(file: bytes = File(...)):

    image = Image.open(io.BytesIO(file)).convert("RGB")
    stats = process_image(image)
    return stats

After starting the server, I can manually test the endpoint successfully by running a call with requests

import requests
from requests_toolbelt.multipart.encoder import MultipartEncoder

url = "http://127.0.0.1:8000/analyse"

filename = "./example.jpg"
m = MultipartEncoder(
        fields={'file': ('filename', open(filename, 'rb'), 'image/jpeg')}
    )
r = requests.post(url, data=m, headers={'Content-Type': m.content_type}, timeout = 8000)
assert r.status_code == 200

However, setting up tests in a function of the form:

from fastapi.testclient import TestClient
from requests_toolbelt.multipart.encoder import MultipartEncoder
from app.server import app

client = TestClient(app)

def test_image_analysis():

    filename = "example.jpg"

    m = MultipartEncoder(
        fields={'file': ('filename', open(filename, 'rb'), 'image/jpeg')}
        )

    response = client.post("/analyse",
                           data=m,
                           headers={"Content-Type": "multipart/form-data"}
                           )

    assert response.status_code == 200

when running tests with python -m pytest , that gives me back a

>       assert response.status_code == 200
E       assert 400 == 200
E        +  where 400 = <Response [400]>.status_code

tests\test_server.py:22: AssertionError
-------------------------------------------------------- Captured log call --------------------------------------------------------- 
ERROR    fastapi:routing.py:133 Error getting request body: can't concat NoneType to bytes
===================================================== short test summary info ====================================================== 
FAILED tests/test_server.py::test_image_analysis - assert 400 == 200

what am I doing wrong?
What's the right way to write a test function test_image_analysis() using an image file?

You see a different behavior because requests and TestClient are not exactly same in every aspect as TestClient wraps requests . To dig deeper, refer to the source code: ( FastAPI is using TestClient from starlette library, FYI)

https://github.com/encode/starlette/blob/master/starlette/testclient.py

To solve, you can get rid of MultipartEncoder because requests can accept file bytes and encode it by form-data format, with something like

# change it
r = requests.post(url, data=m, headers={'Content-Type': m.content_type}, timeout = 8000)

# to 
r = requests.post(url, files={"file": ("filename", open(filename, "rb"), "image/jpeg")})

and modifying the FastAPI test code:

# change
response = client.post("/analyse",
                       data=m,
                       headers={"Content-Type": "multipart/form-data"}
                       )
# to
response = client.post(
    "/analyse", files={"file": ("filename", open(filename, "rb"), "image/jpeg")}
)

Below code is working for me:

** API Structure: **

File: api_routers.py

from fastapi import APIRouter, File, UploadFile, Query
router = APIRouter()
@router.post(path="{{API_PATH}}", tags=["Prediction"])
def prediction(id: str, uploadFile: UploadFile):
    ...
    {{CODE}}
    return response

Testing Code

File: test_api_router.py

import pytest
import os
from fastapi.testclient import TestClient
import.api_routers

client = TestClient(api_routers.router)

def test_prediction(constants):  
    # Use constants if fixture created
    file_path = "{{IMAGE PATH}}"
    if os.path.isfile(file_path):
        _files = {'uploadFile': open(file_path, 'rb')}
        response = client.post('{{API_PATH}}',
                           params={
                               "id": {{ID}}
                           },
                           files=_files
                           )
        assert response.status_code == 200
    else:
        pytest.fail("Scratch file does not exists.")

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM