Welcome to the third lesson in our "Model Serving with FastAPI" course! So far, we've established a strong foundation by building a basic FastAPI application and integrating a Machine Learning model for diamond price prediction. Now, we'll focus on a crucial aspect of production-ready applications: testing.
Testing is a vital part of developing reliable APIs, especially when they serve Machine Learning models that make critical predictions. By the end of this lesson, you'll learn how to create comprehensive tests for your diamond price prediction API using pytest
, ensuring your endpoints function correctly and handle various scenarios appropriately.
Let's build robust tests that give us confidence in our API's functionality!
Before diving into code, let's understand why testing is particularly important for Machine Learning APIs:
- Model behavior validation - We need to ensure predictions are made correctly and consistently.
- Input validation - Our API must properly handle both valid and invalid inputs.
- Error handling - The system should respond appropriately when something goes wrong.
- Reliability - APIs in production need to be dependable under various conditions.
In this lesson, we'll focus on unit tests that verify our API endpoints function correctly in isolation. We'll use mocking techniques to replace the actual ML model with controlled test doubles, allowing us to test the API's logic independently from the model implementation. This approach lets us verify that our FastAPI routes, data validation, and error handling work as expected without being affected by the complexities of the actual ML model. Note that effective API tests should be isolated from external dependencies (like the actual ML model), cover typical usage scenarios, test error conditions to ensure graceful failure, and verify expected responses, including status codes and data formats.
The pytest
framework provides an excellent foundation for testing our API. It offers a simple syntax while providing powerful features like parametrization and fixtures. FastAPI integrates wonderfully with pytest
through its TestClient
class.
Let's start by setting up our test module:
The TestClient
wraps our FastAPI application, allowing us to make requests directly to our endpoints without actually running a server. This client simulates HTTP requests and captures the responses, making it easy to verify that our endpoints behave as expected.
This setup establishes the foundation for all our tests. The client
variable will be used throughout our test functions to interact with our API endpoints, just as a real client would in production. By importing our application instance directly, we're testing the exact same code that would run in a production environment.
Let's begin by testing our basic endpoints. First, we'll test the root endpoint that welcomes users to our API:
This test makes a GET
request to the root path and verifies that the response has a 200 status code (OK) and that the response JSON contains a welcome message.
Next, let's test the health check endpoint:
Health check endpoints are critical for production systems as they allow monitoring tools to verify your service is operating correctly. This test confirms our endpoint returns a 200 status code and includes both a "healthy" status and the API version in its response.
These simple tests follow a common pattern in API testing that you'll use repeatedly: make a request, check the status code, and verify the response content matches expectations. While straightforward, they provide valuable assurance that your API's foundation is solid.
When unit-testing prediction endpoints, we face a challenge: we don't want our tests to depend on actual model files or be affected by model behavior changes.
In pytest, fixtures are special functions that provide reusable test dependencies. They're a powerful way to set up preconditions for your tests, manage test resources, and inject dependencies. Fixtures help create a consistent testing environment and reduce code duplication across tests.
For our ML API testing, we'll use fixtures to implement mocking — creating simplified substitutes for the model that return predefined values. This approach allows us to test our API logic independently from the actual model implementation.
Let's create a fixture to handle our model mocking:
This fixture creates two mock classes:
MockModel
with apredict
method that always returns a fixed prediction value;MockPreprocessor
with atransform
method that returns a fixed feature array.
By using these mock objects, we can test our API's logic without depending on actual model files or behavior. This makes our tests more reliable and faster to run.
With our mocking fixture in place, we can now test our prediction endpoint with realistic but controlled scenarios:
This test demonstrates several powerful testing techniques:
-
Parametrization with
@pytest.mark.parametrize
lets you run the same test with different inputs — you could easily add more test cases to this list. -
Monkey patching with the
monkeypatch
fixture temporarily replaces functions or attributes at runtime. This built-in pytest fixture allows us to modify behavior without changing the actual code. Here, we use it to replace our realget_model
function with a lambda that returns our mock objects, eliminating the dependency on the actual model file during tests. -
After setting up these mocks, we make a
POST
request with our test diamond features and verify that the response contains the expected prediction (1500.0) and returns the original features. This approach lets us test our API logic independently from the actual Machine Learning model.
A robust API should handle invalid inputs gracefully. Let's test how our prediction endpoint responds to various error conditions:
This test uses parametrization brilliantly to examine multiple error scenarios with minimal code duplication:
- A request missing a required field (carat is omitted).
- A request with an invalid data type (string instead of number for carat).
- A request with an invalid value (negative carat).
In each case, we expect a 422 Unprocessable Entity status code, which FastAPI automatically returns when request validation fails. By testing these scenarios, we ensure our API's data validation works correctly and prevents invalid data from reaching our model.
What happens when our model isn't available? This is a critical scenario to handle properly in production environments. Let's test this case:
This test simulates a scenario where our get_model
function fails to load the model and raises an HTTPException
. We use a clean, straightforward approach to simulate this failure. This tests if our dependency injection system properly propagates errors to clients with appropriate status codes.
Why is this important? In production, models might fail to load for many reasons — corrupt files, memory issues, or version incompatibilities. When this happens, your API should provide clear feedback rather than crashing or hanging. This test ensures your error handling system works correctly for this critical failure mode.
Now that we've written a comprehensive test suite, let's explore how to actually run these tests using pytest
. Running tests regularly is essential to catch issues early in your development process.
First, ensure your test files follow the pytest naming convention - they should be named test_*.py
or *_test.py
to be automatically discovered. A common convention is to place these files in a tests
directory for better organization, but note this is not explicitly required by pytest
.
To run all your tests, simply move to the project directory and execute:
This command discovers and runs all test files in your project. You can run specific test files by providing the path:
You can also run pytest programmatically from Python code:
Excellent work! You've now learned how to create a comprehensive test suite for your FastAPI applications. These tests verify that your endpoints function correctly, validate inputs properly, and handle errors gracefully — all essential qualities for a production-ready API.
Now, time for some practice. Keep up the great work!
