A bit more refactoring and test cleanup