def test_wrap_polynomial_features():
# load data
X = fetch_california_housing(as_frame=True).frame
# prepare selectors
tr = PolynomialFeatures()
tr_wrap = SklearnTransformerWrapper(transformer=PolynomialFeatures())
# Test:
# When passing variable list
varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"]
tr_wrap.set_params(variables=varlist)
Xt = pd.DataFrame(tr.fit_transform(X[varlist]),
columns=tr.get_feature_names_out(varlist))
Xw = tr_wrap.fit_transform(X)
pd.testing.assert_frame_equal(Xw, pd.concat([X, Xt], axis=1))
assert Xw.shape[1] == len(X.columns) + len(
tr.get_feature_names_out(varlist))
# when variable list is None
tr_wrap.set_params(variables=None)
Xt = pd.DataFrame(tr.fit_transform(X), columns=tr.get_feature_names_out())
Xw = tr_wrap.fit_transform(X)
pd.testing.assert_frame_equal(Xw, pd.concat([X, Xt], axis=1))
assert Xw.shape[1] == len(X.columns) + len(
tr.get_feature_names_out(X.columns))
