Bug and unit test fixes