Increase code coverage

More tests, 1 mock