How mock a method with parameters and return an iterable in python-CodePudding

I need to test this function with a unit test:

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]

I wrote, inicialy, this test and worked:

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

But in this test I need to load the model. As this is an unit test, I would like to mock the responses, thus load an external model can be disable. I am trying something like this (and a combination of the lines commented in the code):

def test_nlp_extraction_entrada_correta():  
    texts = 'xilitol'

    doc = Mock(name="DOC")
    ents = Mock(name="ENTS", label_='xilitol', text="xilitol")

    doc.ents = [ents]
    
    from nextmock import Mock
    nlp = Mock()
    nlp_mock = Mock()
    nlp.with_args([texts]).returns([doc])
    nlp_mock.pipe = nlp([texts])
    
    # nlp_mock.pipe.with_args([texts]).returns(doc)
    # nlp_mock.pipe = [Mock(return_value=doc)]

    result = nlp_extraction(texts, nlp=nlp_mock)
    assert result == result_reference

But an error always raise, saying that nlp.pipe([texts]) mock object is not iterable. So, I need to mock this part nlp.pipe([texts]) and return the doc object. How I can do this? Something I am missing in the proccess, can someone help me.

CodePudding user response：

As Cpt.Hook said in comments, the solution was achieved using nlp.pipe.return_value = [doc].