如何模拟与 Unittest 的数据库连接，用于连接作为上下文管理器的类？

Question

Here it is class with one of it's methods running when class is initializing:

class StatCollector:
    def __init__(self, poll_stat) -> None:
        self.polls = self.__get_polls()

    def __get_polls(self) -> Dict[str, Poll]:
        with pyodbc.connect(MSSQL_CONNECTION_PARAMS) as cnxn:
            polls = dict()
            cursor = cnxn.cursor()
            query = self.poll_stat.poll_ids_query_getter()
            cursor.execute(query, self.poll_stat.products_block)
            for poll in map(Poll._make, cursor.fetchall()):
                polls[poll.poll_id] = poll
        return polls

I want to test other methods of this class and my first goal is to fill self.polls with initial values with no real connection to db and using __get_polls method. My try:

@patch("pyodbc.connect")
class testStatCollector(unittest.TestCase):
    def test_initial_values_setted(self, mock_connect):
        cursor = MagicMock(name="my_cursor")
        cursor.fetchall.return_value = [("2", "А", "B")]
        cnxn = MagicMock(name="my_cnxn_mfk")
        cnxn.cursor.return_value = cursor
        mock_connect.return_value.__enter__ = cnxn
        self.test_class = PollsStatCollector(IVR)
        self.assertEqual(
            self.test_class.polls, {"2": Poll("2", "A", "B")}
        )
        self.assertIsInstance(self.test_class.period_start_time, datetime)

But self.polls are empty after execution. I got: AssertionError: {} != {'2': Poll(poll_id='2', product='A', products_block='B')}
and I see in debug, that cnxn name = my_cnxn_mfk when __get_polls executing, but then cursor with default name = <MagicMock name='my_cnxn_mfk().cursor()' id='1883689785424'> .
So i guess that i make mistake in this part cnxn.cursor.return_value = cursor , but i dont know how to fix it.

Answer 1

Mistake was here:

mock_connect.return_value.__enter__ = cnxn

Should be replaced with

mock_connect.return_value.__enter__.return_value = cnxn