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.
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:
importpytestfrom..models.todoimportTaskInput,TaskOutputfromdatetimeimportdate,timedeltadeftest_create_realistic_input_model():input_model=TaskInput(name="write blog post",priority=1,due_date=date.today(),done=False)assertinput_model.name=="write blog post"assertinput_model.priority==1assertinput_model.due_date==date.today()assertinput_model.done==Falsedeftest_check_input_has_a_name_larger_than_five():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="1234",priority=1,due_date=date.today(),done=False)assert"String should have at least 5 characters"instr(e_info.value)deftest_check_input_has_a_name_not_larger_than_100():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="x"*101,priority=1,due_date=date.today(),done=False)assert"String should have at most 100 characters"instr(e_info.value)deftest_check_input_has_a_priority_larger_than_zero():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="write blog post",priority=-1,due_date=date.today(),done=False)assert"Input should be greater than 0"instr(e_info.value)deftest_check_input_has_a_priority_smaler_than_10():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="write blog post",priority=10,due_date=date.today(),done=False)assert"Input should be less than 10"instr(e_info.value)deftest_check_input_does_not_have_a_due_date_in_the_past():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="write blog post",priority=2,due_date=date.today()+timedelta(days=-1),done=False)assertf"Input should be greater than or equal to {date.today()}"instr(e_info.value)deftest_check_input_does_not_have_a_due_date_more_than_a_year_in_the_future():withpytest.raises(ValueError)ase_info:input_model=TaskInput(name="write blog post",priority=2,due_date=date.today()+timedelta(days=+367),done=False)assertf"Input should be less than or equal to {date.today()+timedelta(days=+365)}"instr(e_info.value)
Refactor the models
For our Pydantic models we can add Fields and specify our validation rules:
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:
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:
classTaskInput(BaseModel):name:str=Field(str,min_length=5,max_length=100)priority:int=Field(gt=0,lt=10)due_date:datedone:bool@field_validator('due_date')defdue_date_must_be_between_today_and_one_year_in_the_future(cls,v):ifnotdate.today()<=v<=date.today()+timedelta(days=365):raiseValueError("due_date must be between today and one year in the future")returnv
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:
assertf"due_date must be between today and one year in the future"instr(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.