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