35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
import pandas as pd
|
|
import pytest
|
|
|
|
from delta_barth.analysis import parse
|
|
from delta_barth.errors import FeaturesMissingError
|
|
|
|
|
|
def test_check_needed_features():
|
|
target_features = set(("feat1", "feat2", "feat3"))
|
|
data = pd.DataFrame(
|
|
data=[[1, 2, 3, 4, 5]], columns=["feat1", "feat2", "feat3", "feat4", "feat5"]
|
|
)
|
|
parse.check_needed_features(data, target_features)
|
|
data = pd.DataFrame(
|
|
data=[[1, 2, 3, 4, 5]], columns=["featX", "feat2", "feat3", "feat4", "feat5"]
|
|
)
|
|
with pytest.raises(FeaturesMissingError):
|
|
parse.check_needed_features(data, target_features)
|
|
|
|
|
|
def test_map_features_to_targets():
|
|
feature_map = dict(feat1="feat10", feat2="feat20", feat5="feat50")
|
|
data = pd.DataFrame(
|
|
data=[[1, 2, 3, 4, 5]], columns=["feat1", "feat2", "feat3", "feat4", "feat5"]
|
|
)
|
|
data = parse.map_features_to_targets(data, feature_map)
|
|
assert "feat10" in data.columns
|
|
assert "feat20" in data.columns
|
|
assert "feat50" in data.columns
|
|
assert "feat3" in data.columns
|
|
assert "feat4" in data.columns
|
|
assert "feat1" not in data.columns
|
|
assert "feat2" not in data.columns
|
|
assert "feat5" not in data.columns
|