Adjust tests for new model