Fix unit tests

Removes the mocking, as we were no longer mocking the right thing. It
proves easier to assert that the axis range can be set than to try and
mock a collection of other things.