Fix bug and test