Fix context not being set in restoring of specific mock