Bug fix and adds a unittest.