Fixing test and coverage !!!