Bug fix and test fix