(refactor) Refactor, add helpers, add tests