Mock数据库连接：get_data上下文管理器致数据匹配失败求助

我有一个DatabaseConnector类，在构造函数（__post_init__）中初始化数据库连接。已经成功Mock了数据库连接，检查数据类型时断言能通过，但返回结果值和Mock数据不匹配，导致断言失败——Mock数据返回的是Mock对象。我确定问题出在get_data方法里的上下文管理器，但不知道怎么解决，求指导如何让值匹配。

文件 database_connector.py #

from dataclasses import dataclass
import pg8000
from pg8000 import Connection

# 假设host、user、password、port为已定义的变量
host = "xxx"
user = "xxx"
password = "xxx"
port = 5432

@dataclass
class DatabaseConnector:
    database_engine: str
    instance: str = "read"

    def __post_init__(self):
        self.database_engine = "".join(filter(str.isalnum, self.database_engine.lower()))
        self.instance = "".join(filter(str.isalnum, self.instance.lower()))
        self.connection = self.get_connection()

    def get_connection(self) -> Connection:
        kwargs = {
            "database": host,
            "host": host,
            "user": user,
            "password": password,
            "port": port
        }

        return pg8000.connect(**kwargs)
    
    def get_data(self, query: str) -> list[tuple]:
        with self.connection.cursor() as cursor:
            cursor.execute(query)
            data = cursor.fetchall()

        return data

文件 mock_db.py #

import unittest
from unittest.mock import patch, Mock
# 假设InboundVolume为已定义的类
from your_module import InboundVolume

db_path = "connectors.database.database_connector.DatabaseConnector"

class TestSPO(unittest.TestCase):
    @patch("database_connector.DatabaseConnector.get_data")
    @patch("database_connector.DatabaseConnector.get_connection")
    def test_connection(self, mock_connection, mock_data):
        mock_record = (0, "1")
        mock_data.return_value = (0, "1")
        mock_connection.return_value = Mock()
        expected_result = InboundVolume().get_inbound_volume()
        self.assertEqual(expected_result, mock_record)

问题核心是没正确Mockcursor对象的行为——你现在只Mock了get_connection返回一个空Mock，但这个Mock连接的cursor()返回的还是默认Mock，它的fetchall()不会返回你要的值；而且你还Patch了get_data，相当于直接跳过了get_data的实际逻辑，大概率和InboundVolume.get_inbound_volume()里的调用逻辑不匹配。

按下面的步骤改测试代码就行：

1. 删掉对get_data的Patch，我们要让InboundVolume调用真实的get_data，只Mock底层数据库操作就行。
2. 给Mock的connection配置cursor方法的返回值：创建一个Mock的cursor对象，让connection.cursor()返回它。
3. 给这个Mockcursor的fetchall()方法设置返回值——注意get_data返回的是list[tuple]，所以要把你的mock_record放到列表里。
4. Mock对象默认支持上下文管理器，不用额外配置__enter__，它会自动返回自身。

修改后的测试代码：

import unittest
from unittest.mock import patch, Mock
from your_module import InboundVolume

class TestSPO(unittest.TestCase):
    @patch("database_connector.DatabaseConnector.get_connection")
    def test_connection(self, mock_connection):
        # 定义预期的数据，注意get_data返回的是列表套元组
        mock_records = [(0, "1")]
        # 创建Mock游标，设置fetchall返回预期数据
        mock_cursor = Mock()
        mock_cursor.fetchall.return_value = mock_records
        # 让Mock连接的cursor方法返回这个游标
        mock_connection.return_value.cursor.return_value = mock_cursor
        
        # 执行测试逻辑
        result = InboundVolume().get_inbound_volume()
        
        # 根据InboundVolume的实际返回值调整断言：
        # 如果它返回单个元组，就断言和mock_records[0]相等
        self.assertEqual(result, mock_records[0])
        # 如果它返回整个列表，就用self.assertEqual(result, mock_records)

额外提醒：

• 别同时Patchget_data和get_connection，除非你明确要跳过get_data的逻辑，不然测试就没意义了——Mock最底层的数据库连接才是正确的做法。
• 注意get_data的返回类型是list[tuple]，别直接把单个元组设为fetchall()的返回值，不然类型会不匹配。

内容的提问来源于stack exchange，提问作者PandasM