Add functional tests
Refactor, add classes