Added functional tests and refactoring