如何在 Python 中对包含数据库调用的方法进行单元测试

Question

我想对包含数据库（SQL 服务器）调用的方法进行单元测试。

我不希望测试连接到实际数据库。

我使用 unittest 进行测试，我做了一些研究，似乎 Mocking 可以解决问题，但不确定语法。

下面的代码select声明returns一些整数。我想模拟将针对代码的“cursor.execute”和“cursor.fetchall()”部分。

from databaselibrary.Db import Db

class RandomClass():

    def __init__(self, database):
        self.database = database # Main DB for inserting data
            
    def check_file_status(self, trimmed_file_data, file_date):
        cursor = self.database.cursor()

        cursor.execute(f"""SELECT DISTINCT query_id
                FROM wordcloud_count
                WHERE date = '{file_date}'""")
        
        queries_in_DB = set(row.query_id for row in cursor.fetchall())

        queries_in_file = set(trimmed_file_data.keys())

        if queries_in_DB == queries_in_file:
            return False
        
        return True
    
    def run(self):
        print("Hello")


if __name__ == "__main__":
    
    connection_string = 'sql://user:password@server/database'
    
    database = Db(connection_string, autocommit=True)
    
    random = RandomClass(database)
    
    random.run()

测试 class 可能看起来像这样：

import unittest
from unittest.mock import Mock, patch
from project.RandomClass import RandomClass
from datetime import datetime

class testRandomClass(unittest.TestCase):
    
    def setUp(self):
        self.test_class = RandomClass("don't want to put actual database here")
    
    @patch("project.RandomClass.check_file_status",return_value={123, 1234})
    def test_check_file_status(self):
        
        keys = {'1234':'2','123':'1','111':'5'}
        
        result = self.test_class.check_file_status(keys, datetime(1900, 1, 1, 23, 59, 59))
        
        self.assertTrue(result)

Answer 1

您应该模拟数据库连接对象和游标。然后，将光标的 return 值设置为 return 预期值。我已经测试了下面的代码，并使用 class Row 来模拟来自 fetchall 调用的 return 行：

import unittest
from unittest.mock import MagicMock
from datetime import datetime
from project.RandomClass import RandomClass


class Row(object):
    def __init__(self, x):
        self.query_id = x


class testRandomClass(unittest.TestCase):

    def setUp(self):
        dbc = MagicMock(name="dbconn")
        cursor = MagicMock(name="cursor")
        cursor.fetchall.return_value = [Row(1), Row(2)]
        dbc.cursor.return_value = cursor
        self.test_class = RandomClass(dbc)

    def test_check_file_status(self):
        keys = {'1234': '2', '123': '1', '111': '5'}

        result = self.test_class.check_file_status(keys, datetime(1900, 1, 1, 23, 59, 59))

        self.assertTrue(result)

由于在您的 RandomClass 中您迭代行并获得它们的 query_id，您需要使用 class（或命名元组）作为由模拟编辑的行对象 return .

您应该创建您期望的行对象，并将它们设置为 fetchall 的 return 值。

如何在 Python 中对包含数据库调用的方法进行单元测试

How to unit test a method that contains a database call in Python

python

sql-server

python-3.x

python-unittest

python-unittest.mock