SQLAlchemy单元测试Mock Result异常：无法返回预期二维列表

问题场景 #

编写SQLAlchemy查询的单元测试时，已Mock create_engine、connection、execute等相关方法，但.all()本该返回二维列表，实际却持续返回Mock对象。调试发现必须调用start_date[0][0].return_value才能拿到预期的datetime对象，而非直接通过start_date[0][0]获取。

原测试代码 #

import unittest
import datetime
from unittest.mock import patch, MagicMock
from costing.modules._common import db_query as dbq
from sqlalchemy.engine.base import Engine, Connection
from sqlalchemy.engine.result import Result
from costing.modules._common.db_url import DatabaseClient


class TestDbQuery(unittest.TestCase):
    @patch.dict('os.environ', {
        'PSQL_URL': 'localhost:5432/testdb',
        'PSQL_USER': 'testuser',
        'PSQL_PASS': 'testpass'
    })
    @patch.object(dbq, 'create_engine')
    @patch.object(Engine, 'begin')
    @patch.object(Connection, 'execute')
    @patch.object(Result, 'all')
    def test_query_date_range(self, mock_result, mock_execute, mock_begin, mock_create_engine):
        # Arrange
        start_date = datetime.datetime(2021, 1, 1)
        end_date = datetime.datetime(2021, 1, 31)
        dc = DatabaseClient()
        mock_test = mock_create_engine().begin().__enter__().execute().all()
        mock_test.__getitem__(0).__getitem__(0).return_value = start_date
        mock_result.side_effect = [[start_date], [end_date]]  # Not used during debugging
        mock_create_engine.begin.return_value = mock_begin
        mock_begin.execute.return_value = mock_execute
        mock_execute.return_value.all.return_value = mock_test

        de = dbq.DatabaseEngine(db_url=dc.psql_url)
        result_return = [
            {
                'subscription_id': '123',
                'start_date': '2021-01-01',
                'end_date': '2021-01-31'
            }
        ]

        # Act
        result = de.query_date_range(subscription_id='123')

        # Assert
        
        self.assertIsNotNone(result)
        assert result[0] == result_return[0] 

被测试代码片段 #

with self.engine.begin() as conn:
            start_date: list = conn.execute(start_query).all()
            end_date: list = conn.execute(end_query).all()

        results = [
            {
                'subscription_id': subscription_id,
                'start_date': dt.strftime(start_date[0][0], date_format),
                'end_date': dt.strftime(end_date[0][0], date_format)
            }
        ]

问题根源 #

1. Mock层级设置错误：原代码试图通过mock_test.__getitem__(0).__getitem__(0).return_value = start_date模拟列表取值，但被测试代码中是直接访问start_date[0][0]（索引访问），而非调用方法，这种设置方式只会让start_date[0][0]返回Mock对象，而非实际datetime值。
2. 重复Mock覆盖：同时使用@patch.object(Result, 'all')和手动设置mock_execute.return_value.all.return_value，导致Mock逻辑冲突。
3. 链式Mock逻辑混乱：通过mock_create_engine().begin().__enter__().execute().all()生成的Mock对象，后续赋值未正确关联到实际调用链。

修正后的测试代码 #

import unittest
import datetime
from unittest.mock import patch, MagicMock
from costing.modules._common import db_query as dbq
from sqlalchemy.engine.base import Engine, Connection
from costing.modules._common.db_url import DatabaseClient


class TestDbQuery(unittest.TestCase):
    @patch.dict('os.environ', {
        'PSQL_URL': 'localhost:5432/testdb',
        'PSQL_USER': 'testuser',
        'PSQL_PASS': 'testpass'
    })
    @patch.object(dbq, 'create_engine')
    def test_query_date_range(self, mock_create_engine):
        # Arrange
        start_date_val = datetime.datetime(2021, 1, 1)
        end_date_val = datetime.datetime(2021, 1, 31)
        
        # 构建Mock调用链
        mock_engine = MagicMock(spec=Engine)
        mock_create_engine.return_value = mock_engine
        
        # 处理with语句的上下文返回connection
        mock_begin_ctx = MagicMock()
        mock_conn = MagicMock(spec=Connection)
        mock_begin_ctx.__enter__.return_value = mock_conn
        mock_engine.begin.return_value = mock_begin_ctx
        
        # 给两次execute调用分别设置返回结果
        mock_start_result = MagicMock()
        mock_start_result.all.return_value = [[start_date_val]]
        mock_end_result = MagicMock()
        mock_end_result.all.return_value = [[end_date_val]]
        mock_conn.execute.side_effect = [mock_start_result, mock_end_result]
        
        dc = DatabaseClient()
        de = dbq.DatabaseEngine(db_url=dc.psql_url)
        
        result_return = [
            {
                'subscription_id': '123',
                'start_date': '2021-01-01',
                'end_date': '2021-01-31'
            }
        ]

        # Act
        result = de.query_date_range(subscription_id='123')

        # Assert
        self.assertIsNotNone(result)
        self.assertEqual(result[0], result_return[0])

关键修改说明 #

• 简化Mock层级：移除不必要的patch装饰器，手动构建Mock链，避免层级和顺序冲突。
• 直接匹配返回结构：给result.all()设置二维列表返回值，完全匹配被测试代码中start_date[0][0]的取值逻辑。
• 处理多次调用：通过side_effect给conn.execute设置两次返回结果，对应两次查询的输出。
• 规范断言：使用unittest自带的self.assertEqual替代原生assert，保持测试风格统一。

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