[WIP] Refactoring and tests