import pandas as pd
from tidyfit.preprocessing import (
    fit_standard_scaler,
    apply_standard_scaler,
    fit_one_hot,
    apply_one_hot,
)


def test_scaler_fit_apply():
    df = pd.DataFrame({"a": [1, 2, 3], "b": [2, 4, 6]})
    state = fit_standard_scaler(df, ["a", "b"])
    out = apply_standard_scaler(df, state)
    assert "a" in out.columns and abs(out["a"].mean()) < 1e-9


def test_one_hot_fit_apply():
    df = pd.DataFrame({"c": ["x", "y", "x"]})
    vocab = fit_one_hot(df, ["c"])
    out = apply_one_hot(df, vocab)
    assert "c__x" in out.columns and "c" not in out.columns
