Fixing test issues