Revert "Revert "refactor code & add integration tests.""

This reverts commit 676b60b5e7b3cc76edcb7f7bd454f11de09ecdf7.