Unit Test fix