tests/unit/test_train_split.py

import pandas as pd

from training.train import _xy, time_respecting_split


def test_time_respecting_split_never_lets_test_precede_train():
    df = pd.DataFrame({"step": list(range(1, 101)), "isFraud": [0] * 100})
    train, test = time_respecting_split(df)
    assert train["step"].max() <= test["step"].min()
    assert len(train) + len(test) == len(df)


def test_xy_casts_to_float32_and_int8():
    df = pd.DataFrame({"amount": [1.5, 2.5], "isFraud": [0, 1]})
    X, y = _xy(df, ["amount"])
    assert str(X["amount"].dtype) == "float32"
    assert str(y.dtype) == "int8"