[refactor] Rewriting code to make i testable