Updated tests for new models