refactor code
refactor code and add test