ML API testing with pytest

Testing machine learning APIs for correctness is valuable but often cumbersome. Developers typically need to spin up a server manually and configure it to work with testing frameworks like pytest. LitServe streamlines this process by offering built-in utility methods tailored for pytest, enabling seamless integration without the need for manual server setup. This reduces boilerplate code, accelerates test development, and encourages best practices in model validation by making it easier and faster to test APIs.

Overview

What is pytest

Pytest is a popular Python testing framework. It allows developers to write tests as plain functions using simple assert statements.

Pytest will discover all functions prefixed with test_, run them, and report any failures with helpful error messages.

# A simple function to test
def square(x):
    return x * x

# A test function
def test_square():
    assert square(3) == 9
    assert square(-2) == 4

To run the test, just type: pytest test_math.py

Pytest will run the test_square function and tell us if all assertions passed. If something fails, it shows a helpful error message so we can debug easily.

ML API testing

Testing machine learning APIs is essential to ensure that your models produce reliable, consistent, and correct results when exposed over a network. Unlike typical software functions, ML APIs often involve components like data preprocessing, model inference, and postprocessing—each of which can introduce subtle bugs or inconsistencies.

A solid testing strategy includes:

• Unit tests for individual functions like data cleaning or output formatting.
• Integration tests for end-to-end validation of the API, from request to prediction to response.
• Edge case testing to check how the API handles invalid input, missing fields, or unexpected data types.
• Performance and regression tests to catch slowdowns or accuracy drops after model updates.

Why test ML APIs

ML APIs often combine multiple steps—preprocessing, inference, postprocessing—and are prone to subtle bugs. Testing helps:

1. Validate correctness: Does the model return the right prediction for a given input?
2. Catch regressions: Ensure new changes don’t break existing behavior.
3. Handle edge cases: What happens if required fields are missing or malformed?
4. Improve confidence: Make changes without fear of breaking things.
5. Save time: Prevent costly manual verification by automating checks.

Testing with Pytest and LitServe

LitServe makes testing APIs just as easy as serving them. With the built-in wrap_litserve_start() helper, we can simulate realistic startup and request flows in-memory using TestClient from FastAPI.

Requirements

Install the necessary packages:

pip install pytest httpx asgi_lifespan

Implement a LitServe API

Let's start with the simplest possible LitServe API which can be modified for any model server.

from fastapi.testclient import TestClient
from fastapi import Request, Response
from litserve import LitAPI, LitServer


class SimpleLitAPI(LitAPI):
    def setup(self, device):
        self.model = lambda x: x**2

    def decode_request(self, request: Request):
        return request["input"]

    def predict(self, x):
        return self.model(x)

    def encode_response(self, output) -> Response:
        return {"output": output}

Write test for LitServe API

Now we can test this API without manually starting a server. Create a file named test_api.py:

import pytest, os
from fastapi.testclient import TestClient
from litserve.utils import wrap_litserve_start

def test_simple_prediction_success():
    api = SimpleLitAPI()
    server = ls.LitServer(api)

    with wrap_litserve_start(server) as server, TestClient(server.app) as client:
        response = client.post("/predict", json={"input": 4})
        assert response.status_code == 200
        assert response.json() == {"output": 16}

Key Details:

• wrap_litserve_start(server) ensures setup() and the internal FastAPI lifespan hooks are properly triggered—just like when running the server.
• TestClient(server.app) provides an in-memory client for sending simulated HTTP requests.
• client.post() lets us test /predict just as a real client would.

To run the test:

pytest test_api.py

We should see output similar to this, indicating that our tests passed (or showed the expected errors for the invalid input tests):

============================= test session starts ==============================
...
test_api.py .                                                           [100%]
============================== 1 passed in 0.XXs ===============================

We've successfully written and run our first Pytest tests for a LitServe API.

Advanced

PyTorch model

Let’s test a real torch.nn.Linear model on both CPU and GPU.

import pytest, torch, torch.nn as nn
from fastapi import Request, Response
from fastapi.testclient import TestClient
from litserve import LitAPI, LitServer
from litserve.utils import wrap_litserve_start

class Linear(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
        self.linear.weight.data.fill_(2.0)
        self.linear.bias.data.fill_(1.0)

    def forward(self, x):
        return self.linear(x)

class TorchLitAPI(LitAPI):
    def setup(self, device):
        self.model = Linear().to(device)
        self.device = device

    def decode_request(self, request: Request):
        value = request["input"]
        return torch.tensor([value], device=self.device)

    def predict(self, x):
        return self.model(x[None, :])

    def encode_response(self, output) -> Response:
        return {"output": float(output)}

Write tests for CPU and GPU:

def test_torch_cpu():
    server = LitServer(TorchLitAPI(), accelerator="cpu", devices=1)
    with wrap_litserve_start(server) as server, TestClient(server.app) as client:
        response = client.post("/predict", json={"input": 4.0})
        assert response.status_code == 200
        assert response.json() == {"output": 9.0}

@pytest.mark.skipif(torch.cuda.device_count() == 0, reason="requires CUDA")
def test_torch_gpu():
    server = LitServer(TorchLitAPI(), accelerator="cuda", devices=1)
    with wrap_litserve_start(server) as server, TestClient(server.app) as client:
        response = client.post("/predict", json={"input": 4.0})
        assert response.status_code == 200
        assert response.json() == {"output": 9.0}

Key concepts

• wrap_litserve_start() handles internal FastAPI app startup and teardown.
• TestClient() lets us send requests without an actual network server.
• Torch models can be tested directly using accelerator="cpu" or "cuda".
• pytest.mark.skipif() ensures GPU tests only run on machines with CUDA available.

Conclusion

LitServe makes ML API testing fast, production-realistic, and easy to automate:

✅ No manual server launch

✅ Supports CPU and GPU inference

✅ Pytest-compatible with full lifecycle support

✅ Great for both local dev and CI pipelines

With a few lines of code, we can verify the full stack: request handling, model prediction, and response formatting.