Fix undefined method
Refactoring