如何正确Mock OracleDB连接生成的Cursor的execute方法？

问题 #

需要为以下Python代码编写单元测试：

import os
import oracledb
import logging

logger = logging.getLogger(__name__)

class User:
    def __init__(self, first_name, last_name, mail):
        self.first_name = first_name
        self.last_name = last_name
        self.mail = mail

class DatabaseCaller:

    def __init__(self):
        self.host = os.getenv("HOSTNAME")
        self.port = os.getenv("PORT")
        self.service_name = os.getenv("SID")
        self.user = os.getenv("USER")
        self.password = os.getenv("PASSWORD")


    def user_exists(self, my_user: User) -> bool:
        this_user_exists: bool = True
        first_name_lowered = my_user.first_name.lower()
        last_name_lowered = my_user.last_name.lower()
        mail_lowered = my_user.mail.lower()
        sql_request = """SELECT email
            FROM users
            WHERE lower(firstname) = :first_name 
            AND lower(lastname) = :last_name 
            AND lower(emailaddress) = :mail"""
        bind_named_parameters = dict(first_name=first_name_lowered, last_name=last_name_lowered, mail=mail_lowered)
        try:
            with oracledb.connect(user=self.user, password=self.password,
                                  host=self.host, port=self.port, service_name=self.service_name) as connection:
                with connection.cursor() as cursor:
                    cursor.execute(sql_request, bind_named_parameters)
                    logger.info(f"Nb row users retrieved {cursor.rowcount}")
                    if cursor.rowcount == 0:
                        this_user_exists = False
        except Exception as exception:
            logger.info(f"Exception launched during the user exists sql method {exception}")
            raise exception
        return this_user_exists

测试时已Mock oracledb.connect，但无法正确模拟连接生成的cursor对象，cursor.execute.assert_called_once()断言始终失败，现有测试代码如下：

import unittest
from unittest.mock import patch, MagicMock
import os
from your_module import DatabaseCaller, User  # 替换为实际模块名

class DBTestCase(unittest.TestCase):

    def test_user_exists(self):
        with patch("oracledb.connect") as mock_connection:
            cursor = mock_connection.cursor.return_value
            cursor_bis = mock_connection.return_value.__enter__.cursor.return_value.__enter__
            print(type(mock_connection.return_value.__enter__.cursor.return_value.__enter__))
            print(type(mock_connection.__enter__.cursor.return_value.__enter__))

            mosaic_user = User("Tom", "Rob", "tom.rob@goog.com")

            os.environ["HOSTNAME"] = "url_database"
            os.environ["PORT"] = "99"
            os.environ["SID"] = "sid"
            os.environ["USER"] = "user_for_database"
            os.environ["PASSWORD"] = "xxx"

            real_database_caller = DatabaseCaller()
            real_database_caller.user_exists(mosaic_user)

            assert real_database_caller.host == os.environ["HOSTNAME"]
            assert real_database_caller.port == os.environ["PORT"]
            assert real_database_caller.service_name == os.environ["SID"]
            assert real_database_caller.user == os.environ["USER"]
            assert real_database_caller.password == os.environ["PASSWORD"]

            mock_connection.assert_called_once_with(
                user=os.environ["USER"],
                password=os.environ["PASSWORD"],
                host=os.environ["HOSTNAME"],
                port=os.environ["PORT"],
                service_name=os.environ["SID"]
            )
            
            # 断言失败，cursor.execute未被识别为调用
            cursor.execute.assert_called_once()

请问如何正确配置Mock，以验证cursor.execute方法的调用？

解决方案 #

问题出在没有正确模拟嵌套上下文管理器的层级：oracledb.connect返回的连接对象是上下文管理器，connection.cursor()返回的cursor也是上下文管理器，需要逐层Mock这些对象的__enter__返回值。

正确的Mock步骤： #

1. 模拟oracledb.connect的上下文返回值（即实际的数据库连接对象）
2. 模拟该连接对象的cursor()方法，返回cursor上下文管理器
3. 模拟cursor上下文管理器的__enter__返回值，即实际用于执行SQL的cursor实例

修改后的测试代码： #

import unittest
from unittest.mock import patch, MagicMock
import os
from your_module import DatabaseCaller, User  # 替换为实际模块名

class DBTestCase(unittest.TestCase):

    def test_user_exists(self):
        with patch("oracledb.connect") as mock_connect:
            # 1. 模拟connect上下文返回的连接对象
            mock_connection = MagicMock()
            mock_connect.return_value.__enter__.return_value = mock_connection

            # 2. 模拟connection.cursor()返回的cursor上下文管理器
            mock_cursor_context = MagicMock()
            mock_connection.cursor.return_value = mock_cursor_context

            # 3. 模拟cursor上下文返回的实际cursor实例
            mock_cursor = MagicMock()
            mock_cursor_context.__enter__.return_value = mock_cursor
            # 设置cursor.rowcount，可根据测试场景调整值
            mock_cursor.rowcount = 1

            # 准备测试数据和环境变量
            mosaic_user = User("Tom", "Rob", "tom.rob@goog.com")
            os.environ.update({
                "HOSTNAME": "url_database",
                "PORT": "99",
                "SID": "sid",
                "USER": "user_for_database",
                "PASSWORD": "xxx"
            })

            # 执行测试方法
            db_caller = DatabaseCaller()
            result = db_caller.user_exists(mosaic_user)

            # 验证环境变量读取正确
            self.assertEqual(db_caller.host, os.environ["HOSTNAME"])
            self.assertEqual(db_caller.port, os.environ["PORT"])
            self.assertEqual(db_caller.service_name, os.environ["SID"])
            self.assertEqual(db_caller.user, os.environ["USER"])
            self.assertEqual(db_caller.password, os.environ["PASSWORD"])

            # 验证connect被正确调用
            mock_connect.assert_called_once_with(
                user=os.environ["USER"],
                password=os.environ["PASSWORD"],
                host=os.environ["HOSTNAME"],
                port=os.environ["PORT"],
                service_name=os.environ["SID"]
            )

            # 验证cursor.execute被调用，还可以进一步验证参数
            mock_cursor.execute.assert_called_once()
            # 可选：验证execute的SQL语句和绑定参数是否正确
            expected_sql = """SELECT email
            FROM users
            WHERE lower(firstname) = :first_name 
            AND lower(lastname) = :last_name 
            AND lower(emailaddress) = :mail"""
            expected_params = {
                "first_name": "tom",
                "last_name": "rob",
                "mail": "tom.rob@goog.com"
            }
            mock_cursor.execute.assert_called_once_with(expected_sql, expected_params)

            # 可选：验证返回值是否符合预期（这里rowcount=1，所以返回True）
            self.assertEqual(result, True)

关键说明： #

• mock_connect.return_value.__enter__.return_value对应with oracledb.connect(...) as connection中的connection对象
• mock_connection.cursor.return_value.__enter__.return_value对应with connection.cursor() as cursor中的cursor对象
• 可以通过设置mock_cursor.rowcount的值，测试不同分支逻辑（比如设为0时，方法应返回False）

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