Second time pull request

1. Add unit test cases 
2. Slightly modify the structure