python - 为同一个函数python模拟两种方法

Question

我想模拟 sklearn 模型的两种方法（predict_proba 和 classes_ ）。我有一个接收模板和文本并返回标签和分数的函数。

import numpy as np
from unittest.mock import MagicMock

def model_predict_proba(model, text):
    pred_proba_model = model.predict_proba([text])
    score = pred_proba_model.max()
    label = model.classes_[np.argmax(pred_proba_model)]
    return label, score


def test_model_predict_proba():
    mock_model = MagicMock()
    mock_model.predict_proba.return_value = np.array([0.90, 0.23])
    mock_model.classes_.return_value= np.array(['FOOD', 'DRINK'])
    text = 'Apple pie'

    expected = ("FOOD", 0.90)
    result = model_predict_proba(mock_model, text)
    
    assert result == expected

当我运行此测试时，我收到以下错误消息：

有人能帮我吗？

标签： pythonpython-3.xmockingpytestpython-unittest

Answer 1

这应该可以解决问题：

def model_predict_proba(model, text):
    pred_proba_model = model.predict_proba([text])
    score = pred_proba_model.max()
    label = model.classes_[np.argmax(pred_proba_model)]
    return label, score


def test_model_predict_proba():
    mock_model = MagicMock()
    mock_model.predict_proba.return_value.max.return_value = 0.90
    mock_model.classes_.__getitem__.return_value ='FOOD'
    text = 'Apple pie'

    expected = ("FOOD", 0.90)
    result = model_predict_proba(mock_model, text)

    assert result == expected

请注意，由于您正在模拟您的模型，因此此测试实际上并没有以任何有用的方式测试模型——我假设您编写此函数只是作为了解其MagicMock工作原理的练习。模拟的目的通常是模拟您正在测试的事物的输入或依赖关系，而不是模拟事物本身。