Cleaned up my code after the refactor