Fixing unit test issues