python - 来自导入模块的 Python 模拟函数

Question

如何为以下代码编写单元测试用例。该代码仅返回一个数据库连接。我无法创建真正的数据库连接。所以，我依赖于对象模拟。

文件：code.py

import snowflake.connector
import logging
import simplejson as json


def get_database_connection(user, password, account, warehouse, database=None, schema=None):
    logging.info("DB Connection : START : " + __name__)
    try:
        return snowflake.connector.connect(
            user=user,
            password=password,
            account=account,
            warehouse=warehouse,
            database=database,
            schema=schema
        )
    except Exception:
        logging.exception("Exception while creating connection")

测试代码.py

class SnowflakeDbTest(unittest.TestCase):

    def test_get_database_connection(self):
         snowflake = Mock()
         snowflake.connector.return_value = Mock()
         snowflake.connector.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", code.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))

解决方案

    @patch('code.snowflake.connector')
    def test_get_database_connection(self, mock_connection):
         mock_connection.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", snowflakedb.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))

如何模拟在“code.py”中导入的“snowflake.connector”

Answer 1

修补是解决方案

    @patch('code.snowflake.connector')
    def test_get_database_connection(self, mock_connection):
         mock_connection.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", snowflakedb.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))