Improve unit test coverage

There isn't any at the moment.