Python中如何为MySQL查询类无确定返回值函数编写结构化unittest测试

一、先修正现有代码的结构问题 #

你当前的代码混淆了业务逻辑和测试逻辑，需要先把查询商品的业务函数单独抽离，不要写在unittest的测试类中：

# 单独存放业务逻辑的文件比如product_dao.py
def get_all_products(input=None):
    connection = get_sql_connection()
    cursor = connection.cursor()
    if input:
        # 注意这里原来的写法有SQL注入风险，建议用参数化查询
        query = "SELECT * FROM `online-store`.product WHERE name like %s;"
        cursor.execute(query, (f"%{input}%",))
    else:
        query = "SELECT * FROM `online-store`.product;"
        cursor.execute(query)
    response = cursor.fetchall()
    for row in response:
        print(f"Id={row[0]}\tName={row[1]}\tPrice =£{row[2]}\tSupplier={row[4]}")
    cursor.close()
    connection.close()
    return response

注意：你原来的SQL拼接写法存在SQL注入漏洞，上面改成了参数化查询的正确写法。

二、单元测试的实现思路 #

测试数据库操作类函数时，必须保证测试环境的数据是可控的，所以需要提前在测试库中插入固定的测试数据，再运行断言判断返回结果是否符合预期。

三、完整的测试代码示例 #

import unittest
from product_dao import get_all_products
# 你自己的测试库连接方法，不要用生产库做测试
from your_db_config import get_test_sql_connection

class TestProductQuery(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 所有测试用例运行前执行一次：初始化测试库、插入测试数据
        cls.conn = get_test_sql_connection()
        cursor = cls.conn.cursor()
        # 清空测试表，插入3条固定测试数据
        cursor.execute("TRUNCATE TABLE `online-store`.product;")
        test_products = [
            (1, "苹果手机", 5999, "128G", "苹果供应商"),
            (2, "华为手机", 4999, "256G", "华为供应商"),
            (3, "苹果耳机", 1299, "无线", "苹果供应商")
        ]
        cursor.executemany(
            "INSERT INTO `online-store`.product (id, name, price, spec, supplier) VALUES (%s, %s, %s, %s, %s)",
            test_products
        )
        cls.conn.commit()
        cursor.close()

    @classmethod
    def tearDownClass(cls):
        # 所有测试用例运行完后执行：清空测试表、关闭连接
        cursor = cls.conn.cursor()
        cursor.execute("TRUNCATE TABLE `online-store`.product;")
        cls.conn.commit()
        cursor.close()
        cls.conn.close()

    def test_get_all_products_no_input(self):
        # 测试无搜索参数的场景：应该返回全部3条数据
        result = get_all_products()
        self.assertEqual(len(result), 3)
        # 校验第一条数据的名称和价格是否符合测试数据
        self.assertEqual(result[0][1], "苹果手机")
        self.assertEqual(result[0][2], 5999)

    def test_get_all_products_with_keyword(self):
        # 测试带搜索关键词的场景：搜索“苹果”应该返回2条数据
        result = get_all_products("苹果")
        self.assertEqual(len(result), 2)
        # 校验搜索结果不包含华为手机
        product_names = [row[1] for row in result]
        self.assertNotIn("华为手机", product_names)

if __name__ == '__main__':
    unittest.main()

四、无需连接数据库的Mock测试方案 #

如果不需要实际连接数据库，只测试函数逻辑正确性，可以用mock替换数据库连接操作，适合纯单元测试场景：

import unittest
from unittest.mock import patch, Mock
from product_dao import get_all_products

class TestProductQueryMock(unittest.TestCase):
    @patch('product_dao.get_sql_connection')
    def test_get_all_products_with_keyword(self, mock_conn):
        # 模拟游标返回的固定结果
        mock_cursor = Mock()
        mock_cursor.fetchall.return_value = [
            (1, "苹果手机", 5999, "128G", "苹果供应商"),
            (3, "苹果耳机", 1299, "无线", "苹果供应商")
        ]
        mock_conn.return_value.cursor.return_value = mock_cursor

        result = get_all_products("苹果")
        # 校验SQL语句是否正确生成
        mock_cursor.execute.assert_called_once_with(
            "SELECT * FROM `online-store`.product WHERE name like %s;",
            ("%苹果%",)
        )
        self.assertEqual(len(result), 2)

if __name__ == '__main__':
    unittest.main()

五、测试结果说明 #

运行上述代码后，如果所有用例通过，会输出OK；如果有不符合预期的情况，会明确告知哪条用例失败、预期值和实际值的差异，即可作为符合要求的测试证明。

内容的提问来源于stack exchange，提问作者user17205637