如何使用Pytest测试读取Excel数据的Python函数？

1. 修复UnboundLocalError错误 #

你的函数里first_value和second_value仅在循环找到匹配值时才会被定义，若传入的g在Excel第三列不存在，这两个变量就没有绑定值，返回时会直接抛出UnboundLocalError。解决方法是在函数开头初始化这两个变量：

def function(g):
    first_value = None  # 初始化默认值
    second_value = None
    wb = openpyxl.load_workbook("name_of_execel.xlsx")
    ws = wb.active
    for row in ws.iter_rows(min_col=3, max_col=3): 
        for cell in row:
            if cell.value == g:
                first_value = ws.cell(row=cell.row, column=6).value
                second_value = ws.cell(row=cell.row, column=7).value
                # 找到匹配后立即退出循环，提升效率
                break
        else:
            continue
        break
    return first_value, second_value

2. Pytest测试方案 #

方案一：直接使用真实Excel文件测试 #

Pytest完全可以直接访问Excel文件，无需额外编写辅助读取函数，直接在测试用例中调用你的函数，传入已知测试值并断言结果即可。注意将测试用Excel文件放在测试目录下，或使用绝对路径避免路径问题。

示例测试代码：

import pytest
from your_module import function  # 替换为你的函数所在模块名

def test_function_existing_value():
    # 假设Excel第三列存在"test_g"，对应第6、7列值为"expected_first"和"expected_second"
    first, second = function("test_g")
    assert first == "expected_first"
    assert second == "expected_second"

def test_function_non_existing_value():
    # 测试不存在的g值，应返回初始化的默认值（如None）
    first, second = function("non_existent_g")
    assert first is None
    assert second is None

方案二：用Mock替代Excel文件（推荐） #

若不想依赖真实Excel文件（比如CI环境无文件、需提升测试速度），可使用unittest.mock模拟openpyxl的行为，避免实际读取文件。

示例测试代码：

from unittest.mock import patch, MagicMock
from your_module import function

def test_function_with_mock():
    # 模拟匹配的单元格
    mock_cell = MagicMock()
    mock_cell.value = "test_g"
    mock_cell.row = 5
    
    mock_row = MagicMock()
    mock_row.__iter__.return_value = [mock_cell]
    
    # 模拟工作表
    mock_ws = MagicMock()
    mock_ws.iter_rows.return_value = [mock_row]
    mock_ws.cell.side_effect = lambda row, column: MagicMock(value="val6" if column==6 else "val7")
    
    # 模拟工作簿
    mock_wb = MagicMock()
    mock_wb.active = mock_ws
    
    with patch("openpyxl.load_workbook", return_value=mock_wb):
        first, second = function("test_g")
        assert first == "val6"
        assert second == "val7"
        # 验证load_workbook被正确调用
        openpyxl.load_workbook.assert_called_once_with("name_of_execel.xlsx")

方案三：参数化批量测试 #

若需测试多组输入输出，可先读取Excel中的测试数据，用@pytest.mark.parametrize批量生成测试用例：

import openpyxl
import pytest
from your_module import function

def load_test_data():
    wb = openpyxl.load_workbook("test_data.xlsx")
    ws = wb.active
    test_data = []
    # 假设第1行是表头，第2行起：第3列是输入g，第6列是预期first，第7列是预期second
    for row in ws.iter_rows(min_row=2, values_only=True):
        g = row[2]
        expected_first = row[5]
        expected_second = row[6]
        test_data.append((g, expected_first, expected_second))
    return test_data

@pytest.mark.parametrize("g, expected_first, expected_second", load_test_data())
def test_function_parametrized(g, expected_first, expected_second):
    first, second = function(g)
    assert first == expected_first
    assert second == expected_second

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