Modifications to add test function for models