The basic models we used for our tasks are good to make sure we get the right types. However, if we accept something like int or string, we could get enormously large inputs that are valid strings or numbers but make no sense at all for our application. Let us explore ways to limit our fields to reasonable lengths.
This post is part of my journey to learn Python. You find the code for this post in my PythonFriday repository on GitHub.
Update the existing tests
If we introduce new rules to our models, our existing tests may start to fail. The tests got created before we had those limits in place and while the functionality may stay the same, our tests may create no longer valid models.
We can define the new rules and change the tests to comply with them before we make the changes to the models, or we focus on the tests for the new rules in the models and fix the existing tests later. Both approaches work and are fine. Just make sure that everyone in the team is on the same page and understands what your priorities are.
Define the new rules
To be more strict in what tasks we accept with our API, we can define these rules:
- The name must be between 5 and 100 characters.
- The priority must be greater than 0 and less than 10.
- The due_date can be between today and one year in the future.
Create the new tests
We can create a new test file test_models.py and add the following tests one by one while we make them pass before we add the next one:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import pytest from ..models.todo import TaskInput, TaskOutput from datetime import date, timedelta def test_create_realistic_input_model(): input_model = TaskInput(name="write blog post", priority=1, due_date=date.today(), done=False) assert input_model.name == "write blog post" assert input_model.priority == 1 assert input_model.due_date == date.today() assert input_model.done == False def test_check_input_has_a_name_larger_than_five(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="1234", priority=1, due_date=date.today(), done=False) assert "String should have at least 5 characters" in str(e_info.value) def test_check_input_has_a_name_not_larger_than_100(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="x"*101, priority=1, due_date=date.today(), done=False) assert "String should have at most 100 characters" in str(e_info.value) def test_check_input_has_a_priority_larger_than_zero(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="write blog post", priority=-1, due_date=date.today(), done=False) assert "Input should be greater than 0" in str(e_info.value) def test_check_input_has_a_priority_smaler_than_10(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="write blog post", priority=10, due_date=date.today(), done=False) assert "Input should be less than 10" in str(e_info.value) def test_check_input_does_not_have_a_due_date_in_the_past(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="write blog post", priority=2, due_date=date.today() + timedelta(days=-1), done=False) assert f"Input should be greater than or equal to {date.today()}" in str(e_info.value) def test_check_input_does_not_have_a_due_date_more_than_a_year_in_the_future(): with pytest.raises(ValueError) as e_info: input_model = TaskInput(name="write blog post", priority=2, due_date=date.today() + timedelta(days=+367), done=False) assert f"Input should be less than or equal to {date.today() + timedelta(days=+365)}" in str(e_info.value) |
Refactor the models
For our Pydantic models we can add Fields and specify our validation rules:
|
1 2 3 4 5 6 |
from pydantic import BaseModel, Field class TaskInput(BaseModel): name: str = Field(str, min_length=5, max_length=100) priority: int = Field(gt=0, lt=10) due_date: date = Field(ge=date.today(), le=date.today() + timedelta(days=365)) done: bool |
Fix the OpenAPI exception
While the code above works outside of FastAPI and in our tests, it provokes an exception as soon as we try to open our OpenAPI specification:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
E pydantic_core._pydantic_core.ValidationError: 4 validation errors for OpenAPI E components.schemas.TaskInput.Schema.properties.due_date.Schema.maximum E Input should be a valid number [type=float_type, input_value=datetime.date(2025, 3, 10), input_type=date] E For further information visit https://errors.pydantic.dev/2.6/v/float_type E components.schemas.TaskInput.Schema.properties.due_date.Schema.minimum E Input should be a valid number [type=float_type, input_value=datetime.date(2024, 3, 10), input_type=date] E For further information visit https://errors.pydantic.dev/2.6/v/float_type E components.schemas.TaskInput.Schema.properties.due_date.bool E Input should be a valid boolean [type=bool_type, input_value={'format': 'date', 'maxim...Date', 'type': 'string'}, input_type=dict] E For further information visit https://errors.pydantic.dev/2.6/v/bool_type E components.schemas.TaskInput.Reference.$ref E Field required [type=missing, input_value={'properties': {'name': {...nput', 'type': 'object'}, input_type=dict] E For further information visit https://errors.pydantic.dev/2.6/v/missing ...\Programs\Python\Python312\Lib\site-packages\fastapi\openapi\utils.py:530: ValidationError |
While the exception is nearly useless, the problem comes from the due_date field. The code to create the OpenAPI specification just does not like fields on date values.
To make sure we get notified when such errors occur, we add this test that just fetches data from the OpenAPI endpoint and will pass if there is no exception:
|
1 2 |
def test_docs_endpoint_works(): response = client.get("/openapi.json") |
We now need to change a few places at once to get everything back to a working state. The order does not matter much, but we may not be able to change the code while all tests stay green.
To fix the OpenAPI problem, we need to remove the Field from due_date and add a validator:
|
1 2 3 4 5 6 7 8 9 10 11 |
class TaskInput(BaseModel): name: str = Field(str, min_length=5, max_length=100) priority: int = Field(gt=0, lt=10) due_date: date done: bool @field_validator('due_date') def due_date_must_be_between_today_and_one_year_in_the_future(cls, v): if not date.today() <= v <= date.today() + timedelta(days=365): raise ValueError("due_date must be between today and one year in the future") return v |
Our OpenAPI test now will work, but the two tests who check the valid due_date start to fail. We can fix them by replacing the assert to match our new validation error message:
|
1 |
assert f"due_date must be between today and one year in the future" in str(e_info.value) |
Our tests should now all pass and the docs endpoint should be back online:
Next
With our strict models in place, we can prevent malicious attackers from flooding our API with endless titles or priority values in the billions. There is much more to the topic of security that we will address in a later post. Next week we revisit the status codes for our endpoints and try to improve them.