Test fixes and minor refactoring