delta-barth-py/tests/analysis/test_parse.py

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