Unit test bug fixes