Test code fixes and clean-up