Fixing bugs from refactoring