Refactoring and adding few integration tests