Fix for bugs created in refactoring process