Add unit tests to test existing functionality

[tetigi]

Make sure to include mocking to some degree.

[cliffsun91]

I think we should do some refactoring along the way, i.e. if we find something particularly hard to unit test, then its a sign of bad design in the code, so we should make changes to it. Maybe minimal ones at first.