Python单元测试：如何Mock带参数方法并返回可迭代对象

问题背景 #

需要为以下nlp_extraction函数编写无外部依赖的单元测试：

def nlp_extraction(texts, nlp=None):
    extr = []
    for doc in nlp.pipe([texts]):
        extr.append(list([ent.label_, ent.text]) for ent in doc.ents)
        extracao = [list(extr[i]) for i in range(len(extr))]
        extracao = list(chain.from_iterable(extracao))
        extracao = " ".join([item[1] for item in extracao])
    return [texts, extracao]

最初的测试依赖加载外部NLP模型，虽能运行但不符合单元测试无依赖的要求：

def test_nlp_extraction_entrada_correta():
    nlp = loadModel('ner_extract_ingredients')
    result_reference = ['xilitol', 'xilitol']
    texts = 'xilitol'
    result = nlp_extraction(texts, nlp)
    assert result == result_reference

尝试用Mock替代模型调用时，始终报错「nlp.pipe([texts]) mock对象不可迭代」，当前Mock代码存在方法配置错误。

解决方案 #

核心是正确Mocknlp.pipe方法，让它在被调用时返回可迭代的预设doc对象列表。具体步骤：

• 创建模拟的doc对象，并为其设置ents属性（包含带label_和text属性的模拟实体）
• 配置Mock的nlp对象，使其pipe方法被调用时返回包含预设doc的可迭代对象

正确的测试代码 #

使用标准unittest.mock的实现（若使用nextmock，逻辑一致，仅需调整语法）：

from unittest.mock import Mock
from itertools import chain

def test_nlp_extraction_entrada_correta():  
    texts = 'xilitol'
    result_reference = [texts, 'xilitol']

    # 模拟实体对象，仅需设置函数用到的属性
    mock_ent = Mock()
    mock_ent.label_ = 'INGREDIENT'
    mock_ent.text = 'xilitol'

    # 模拟doc对象，绑定实体列表
    mock_doc = Mock()
    mock_doc.ents = [mock_ent]

    # 模拟nlp对象，配置pipe方法返回可迭代的doc列表
    mock_nlp = Mock()
    mock_nlp.pipe.return_value = [mock_doc]

    # 调用待测试函数并断言结果
    result = nlp_extraction(texts, nlp=mock_nlp)
    assert result == result_reference

关键说明 #

• mock_nlp.pipe.return_value = [mock_doc]是核心配置：让pipe方法被调用时返回包含模拟doc的列表，确保for doc in nlp.pipe([texts])能正常迭代
• 若需要严格匹配pipe的调用参数，可根据Mock库语法调整，比如mock_nlp.pipe.with_args([texts]).return_value = [mock_doc]
• 模拟实体时仅需设置函数实际用到的label_和text属性，无需额外配置其他属性

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