Bug fixes, and fix test