Bug fixed and implement a test class