Further refactoring to simplify and test the code