That makes sense. Since this is a bit of an unusual case, I would spell it in a slightly longer-winded way and make it clear that you’re testing whether the returned function is the same as the argument:
assert (orig := a.call) is f(orig)
or, of course, breaking it out onto two lines if that is appropriate to what you’re doing.