Yi - do refactoring; add mockito