Bug fix and refactor tests