Refactored models for tests. All unit tests passing now.