Fix mock function restoration in test :(

An argument for using decorators, context managers and stuff so this is
impossible to mess up.