Cleaning up unit test, adding mock instead of stub.