fix UT and refacto