def test_FitPredictWithDictionaryOutput__predict(self):
X = self.X
y = self.y.astype(int)
gm = GaussianNB()
wrapped_gm = add_mixins_to_step(gm)
double_wrap = add_mixins_to_step(wrapped_gm)
double_wrap.fit(X=X, y=y)
result = double_wrap.predict_dict(X=X)
self.assertEqual(
sorted(list(result.keys())),
sorted(['predict', 'predict_proba', 'predict_log_proba']))
self.assertEqual(result['predict'].shape[0], self.size)
self.assertEqual(result['predict_proba'].shape[0], self.size)
self.assertEqual(result['predict_log_proba'].shape[0], self.size)
def test_Pipegraph__step__predict_lm(self):
X = self.X
y = self.y
lm = LinearRegression()
lm_step = add_mixins_to_step(lm)
lm_step.pg_fit(X=X, y=y)
assert_array_equal(lm.predict(X), lm_step.pg_predict(X=X)['predict'])
def test_wrap_adaptee_in_process__right_classes(self):
tests_table = [{
'model': LinearRegression(),
'expected_class': FitPredictMixin
}, {
'model': MinMaxScaler(),
'expected_class': FitTransformMixin
}, {
'model': DBSCAN(),
'expected_class': AtomicFitPredictMixin
}, {
'model':
Demultiplexer(),
'expected_class':
CustomFitPredictWithDictionaryOutputMixin
}]
for test_dict in tests_table:
model = test_dict['model']
step = add_mixins_to_step(model)
self.assertEqual(isinstance(step, test_dict['expected_class']),
True)
def test_wrap_adaptee_in_process__wrap_process_does_nothing(self):
lm = LinearRegression()
wrapped_lm = add_mixins_to_step(lm)
double_wrap = add_mixins_to_step(wrapped_lm)
self.assertEqual(double_wrap, lm)
def test_FitPredictWithDictionaryOutput__get_predict_signature(self):
lm = LinearRegression()
wrapped_lm = add_mixins_to_step(lm)
double_wrap = add_mixins_to_step(wrapped_lm)
result = double_wrap._get_predict_signature()
self.assertEqual(result, ['X'])
def test_FitPredictWithDictionaryOutput__get_fit_signature(self):
lm = LinearRegression()
wrapped_lm = add_mixins_to_step(lm)
double_wrap = add_mixins_to_step(wrapped_lm)
result = double_wrap._get_fit_signature()
self.assertEqual(sorted(result), sorted(['X', 'y', 'sample_weight']))