Test fixes and refactoring