This isn't a lot of code, but could it be shortened by just using the `Mock` object from the `mock` library?