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

Question

我想对包含数据库 (SQL Server) 调用的方法进行单元测试。

我不想让测试连接到实际的数据库。

我使用unittest进行测试，我做了一些研究，似乎Mocking可以做到这一点，但不确定语法。

下面代码中的 select 语句返回一些整数。我猜想模拟将针对代码的“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()

测试类可能如下所示：

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

您应该模拟数据库连接对象和光标。然后，设置游标的返回值，使其返回期望的值。我测试了下面的代码，并用于class Row模拟从调用返回的fetchall行：

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，因此您需要使用一个类（或命名元组）作为模拟返回的行对象。

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