refactor code
add assert count test case