Unit test bug fix