Bug fixes and new test