Fixing bug in unitTest