[WIP] refactoring and updating unit tests