Bug fix and test improvements