tests/atomic/test_atomic_features.py

import socceraction.atomic.spadl as spadlcfg
from pandas import testing as tm
from pandera.typing import DataFrame
from socceraction.atomic.spadl import AtomicSPADLSchema
from socceraction.atomic.vaep import features as fs

xfns = [
    fs.actiontype,
    fs.actiontype_onehot,
    fs.bodypart,
    fs.bodypart_detailed,
    fs.bodypart_onehot,
    fs.bodypart_detailed_onehot,
    fs.team,
    fs.time,
    fs.time_delta,
    fs.location,
    fs.polar,
    fs.movement_polar,
    fs.direction,
    fs.goalscore,
]


def test_same_index(atomic_spadl_actions: DataFrame[AtomicSPADLSchema]) -> None:
    """The feature generators should not change the index of the input dataframe."""
    atomic_spadl_actions.index += 10
    game_actions_with_names = spadlcfg.add_names(atomic_spadl_actions)
    gamestates = fs.gamestates(game_actions_with_names, 3)
    gamestates = fs.play_left_to_right(gamestates, 782)
    for fn in xfns:
        features = fn(gamestates)
        tm.assert_index_equal(features.index, atomic_spadl_actions.index)