Fixed some tests, added more mocks, few more slight adjustments to code.