Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_encode_in_pipelin_fit_n_transform():
drop_name = pdp.ColDrop('name')
encode_stage = Encode()
pline = drop_name + encode_stage
df = _some_df()
with pytest.raises(UnfittedPipelineStageError):
res_df = pline.transform(df)
res_df = pline.fit(df)
assert 'lbl' in res_df.columns
assert 'name' in res_df.columns
assert res_df['lbl'][1] == 'acd'
assert res_df['lbl'][2] == 'alk'
assert res_df['lbl'][3] == 'alk'
res_df = pline.transform(df)
def test_attribute_stage():
"""Testing attribute pipeline stages."""
pipeline = pdp.ColDrop('name').Bin({'speed': [5]}, drop=True)
assert isinstance(pipeline, PdPipeline)
assert isinstance(pipeline[0], ColDrop)
assert isinstance(pipeline[1], Bin)
df = _some_df()
res_df = pipeline(df)
assert 'speed' in res_df.columns
assert 'name' not in res_df.columns
def test_encode_in_pipeline():
drop_name = pdp.ColDrop('name')
encode_stage = Encode()
pline = drop_name + encode_stage
df = _some_df()
res_df = pline(df)
assert 'lbl' in res_df.columns
assert 'name' not in res_df.columns
assert res_df['lbl'][1] == 0
assert res_df['lbl'][2] == 1
assert res_df['lbl'][3] == 1
# check fitted pipeline
df2 = _some_df2()
res_df2 = pline(df2)
assert 'lbl' in res_df2.columns
assert res_df2['lbl'][1] == 1