More unit test fixes