How to call assert_called_with on mock class while keeping custom return value?

Suppose I want to test a class B and check that its member function foo correctly instantiates and calls a class A and calls the bar method on it. When I mock out A to do this, how do I check that it was correctly called, while ensuring the method still returns its parameter as expected by B?

Here is my attempt:

class B():
    def __init__(self):
        pass

    def foo(self, arg):
        a = A()
        res = a.bar(arg)
        print(f'foo called with {arg} and returned {res}')
        assert res == arg
        
        return res

class A():
    def bar(self, arg):
        return arg

import unittest
import unittest.mock

class MockA:
    @unittest.mock.create_autospec
    def bar(self, arg):
        return arg

@unittest.mock.patch("__main__.A", new=MockA)
class MyTestCase(unittest.TestCase):

    def test_case_1(self):
        x = B().foo(1)
        A.bar.assert_called_with(1)

    def test_case_2(self):
        x = B().foo(2)
        A.bar.assert_called_with(2)

    def test_case_3(self):
        x = B().foo(3)
        A.bar.assert_called_with(3)

MyTestCase().test_case_1()

This fails with an assertion error, since res is <MagicMock name='mock()' id='140545715265664'>. Why does the autospec decorator to MockA.bar ignore the return value in this way?