Update tests for new model types