Unit test fixes