refactor Unit test and structure