Cleaned up model and unit tests