Fix warlus operator behavior when called by a function (#11041)
In #10758 we introduced the support for the use of the walrus operator in the test cases. There was a case which was not handled that caused a bug report #11028. This PR aims to fix the issue and also to improve how the walrus operator is handled in the AssertionRewriter class.
Closes #11028
