Bug fix post refactoring