Python - 模拟嵌套的数据库调用

Question

我正在尝试在下面的 some_func.py 中为 some_func 函数编写单元测试。在此测试期间，我不想连接到任何数据库，并且我想模拟对 DB 的任何调用。

由于实际的数据库调用有点嵌套，我无法让它工作。在这种情况下如何修补任何数据库交互？

db_module.py

import MySQLdb

from random_module.config.config_loader import config
from random_module.passwd_util import get_creds

class MySQLConn:

    _conn = None

    def __init__(self):
        self._conn = self._get_rds_connection()

    def get_conn(self):
        return self._conn

    def _get_rds_connection(self):
        """
        Returns a conn object after establishing connection with the MySQL database

        :return: obj: conn
        """

        try:

            logger.info("Establishing connection with MySQL")

            username, password = get_creds(config['mysql_dev_creds'])

            connection = MySQLdb.connect(
                host=config['host'],
                user=username,
                passwd=password,
                db=config['db_name'],
                port=int(config['db_port']))

            connection.autocommit = True

        except Exception as err:
            logger.error("Unable to establish connection to MySQL")
            raise ConnectionAbortedError(err)

        if (connection):
            logger.info("Connection to MySQL successful")
            return connection
        else:
            return None

db_mysql = MySQLConn()

some_func.py

from random_module.utils.db_module import db_mysql

def some_func():

    try: 

        db_conn = db_mysql.get_conn()
        db_cursor = db_conn.cursor()

        results = db_cursor.execute("SELECT * FROM some_table")
        results = db_cursor.fetchall()

        result_set = []

        for row in results:

            result_set.insert(i, row['x'])
            result_set.insert(i, row['y'])

    except Exception:
        logger.error("some error")

    return result_set

目录结构-

src
├──pkg
|   ├── common
|      |___ some_func.py
|      |___ __init__.py
|       
|   ├── utils
|      |___ db_module.py
|      |___ __init__.py
|
| __init__.py

Answer 1

你需要模拟这一行：db_conn = db_mysql.get_conn()

get_conn 方法的返回值是你感兴趣的。

from random_module.utils.db_module import db_mysql

@mock.patch.object(db_mysql, 'get_conn')
def test_some_func(self, mock_get):
    mock_conn = mock.MagicMock()
    mock_get.return_value = mock_conn
    mock_cursor = mock.MagicMock()
    mock_conn.cursor.return_value = mock_cursor

    expect = ...
    result = some_func()

    self.assertEqual(expect, result)
    self.assertTrue(mock_cursor.execute.called)

如您所见，设置这些模拟程序非常复杂。那是因为您在函数内部实例化对象。更好的方法是重构代码以注入游标，因为游标是唯一与此函数相关的东西。更好的方法是创建一个数据库夹具来测试函数是否与数据库正确交互。

Answer 2

您应该在some_module.py 中模拟db_mysql，然后断言预期的调用是在some_func() 执行之后进行的。

from unittest TestCase
from unittest.mock import patch
from some_module import some_func

class SomeFuncTest(TestCase):

    @patch('some_module.db_mysql')
    def test_some_func(self, mock_db):
        result_set = some_func()

        mock_db.get_conn.assert_called_once_with()
        mock_cursor = mock_db.get_conn.return_value
        mock_cursor.assert_called_once_with()

        mock_cursor.execute.assert_called_once_with("SELECT * FROM some_table")
        mock_cursor.fetchall.return_value = [
            {'x': 'foo1', 'y': 'bar1'},
            {'x': 'foo2', 'y': 'bar2'}
        ]
        mock_cursor.fetchall.assert_called_once_with()

        self.assertEqual(result_set, ['foo1', 'bar1', 'foo2', 'bar2'])