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'])