Unable to set return_value on patched mock object

Viewed 31

I would like to make a unit test using pytest library where I need to set return_value of a patched mocked object predict method. But currently, even if I set return_value, it returns also a mock object. What would be the correct approach to this? What should be the approach to test for correct output of the fit_model method?

A toy example - in path src.model_training_service.model I have python file bol.py with contents:

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

class bol:
    def __init__(self):
        self.model = None

    def fit_model(self, x, y):
        self.model = RandomForestClassifier()
        print(type(self.model))
        self.model.fit(x,y)
        y_pred = self.model.predict(x) # this returns mock object instead of intended np.array([9,4,5])
        print('printing y_pred')
        print(type(y_pred))
        acc = accuracy_score(y, y_pred)

        result = 1 + acc
        return result

The test is looking like this:

def test_bol(mocker):
    m = mocker.patch('src.model_training_service.models.bol.RandomForestClassifier')
    m.predict.return_value = np.array([9,4,5])

    x = np.array([[1,2,3,4],
                  [5,6,7,8],
                  [3,2,1,1]])
    y = np.array([8, 4, 2])
    obj = bol()
    result = obj.fit_model(x,y)
    expected_accuracy = 2.

    assert result == expected_accuracy

When I try to put the class and the test into one file, it works as inteded:

Contents of one file:

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

class bol2:
    def __init__(self):
        self.model = None

    def fit_model(self, x, y):
        self.model = RandomForestClassifier()
        print(type(self.model))
        self.model.fit(x,y)
        y_pred = self.model.predict(x)
        print('printing y_pred')
        print(type(y_pred))
        acc = accuracy_score(y, y_pred)
        result = 1 + acc
        return result

def test_bol2(mocker):
    m = mocker.patch('sklearn.ensemble.RandomForestClassifier') # but this is probably wrong, because I am not mocking the object in the module
    m.predict.return_value = np.array([9,4,5])

    x = np.array([[1,2,3,4],
                  [5,6,7,8],
                  [3,2,1,1]])
    y = np.array([8, 4, 2])
    obj = bol2()
    expected_accuracy = 2.
    result = obj.fit_model(x,y)

    assert result == expected_accuracy

Thanks

0 Answers
Related