def test_add_lags_wrong_inputs(test_df):
invalid_df = [[1, 2, 3], [4, 5, 6]]
invalid_lags = ["1", "2"]
with pytest.raises(ValueError, match="lags must be a list of type: ?"):
add_lags(test_df, ["X1"], invalid_lags)
with pytest.raises(ValueError, match="X type should be one of: ?"):
add_lags(invalid_df, ["X1"], 1)
def test_add_lags_correct_df(test_df):
expected = pd.DataFrame({
"X1": [1, 2],
"X2": ["178", "154"],
"X1-1": [0, 1]
})
ans = add_lags(test_df, "X1", -1)
assert (ans.columns == expected.columns).all()
assert (ans.values == expected.values).all()
def test_add_lags_correct_df(test_df):
expected = pd.DataFrame({
'X1': [1, 2],
'X2': ['178', '154'],
'X1-1': [0, 1],
})
ans = add_lags(test_df, 'X1', -1)
assert (ans.columns == expected.columns).all()
assert (ans.values == expected.values).all()
def test_add_lags_correct_X(test_X):
expected = np.array([[-4, 2, -2, 3, 0, -6]])
assert (add_lags(test_X, [0, 1], [1, 2]) == expected).all()
