def test_check_model_1(self):
"""
Unit test check model 1
"""
model = lambda: None
model.predict = types.MethodType(self.predict, model)
_case, _classes = check_model(model)
assert _case == 'regression'
assert _classes is None
def test_check_model_2(self):
"""
Unit test check model 2
"""
model = lambda: None
model._classes = np.array([1, 2])
model.predict = types.MethodType(self.predict, model)
model.predict_proba = types.MethodType(self.predict_proba, model)
_case, _classes = check_model(model)
assert _case == 'classification'
self.assertListEqual(_classes, [1, 2])
def check_model(self):
"""
Check if model has a predict_proba method is a one column dataframe of integer or float
and if y_pred index matches x_pred index
Returns
-------
string:
'regression' or 'classification' according to the attributes of the model
"""
_case, _classes = check_model(self.model)
return _case, _classes
