Refactoring. Also broke one test.