Refactoring to clean up