def test_check_dataframe_column_contains():
df = pd.DataFrame(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
columns=['a', 'b', 'c'])
# this should pass
check_dataframe_column_contains(['a', 'b', 'c'], df)
check_dataframe_column_contains(['a'], df)
check_dataframe_column_contains(['a', 'c'], df)
# this should raise exception
with pytest.raises(ValueError) as e:
check_dataframe_column_contains(['required_column_x'], df)
assert str(e.value).startswith("Missing columns: required_column_x")
with pytest.raises(ValueError) as e:
check_dataframe_column_contains(['a', 'b', 'd', 'e'], df)
assert str(e.value).startswith("Missing columns:")
def test_check_dataframe_column_contains():
df = pd.DataFrame(
np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), columns=["a", "b", "c"]
)
# this should pass
check_dataframe_column_contains({"a": "int", "b": "int", "c": "int"}, df)
check_dataframe_column_contains({"a": "int"}, df)
check_dataframe_column_contains({"a": "int", "c": "int"}, df)
# this should raise exception
with pytest.raises(ValueError) as e:
check_dataframe_column_contains({"required_column_x": "int"}, df)
assert str(e.value).startswith("Missing columns: required_column_x")
with pytest.raises(ValueError) as e:
check_dataframe_column_contains(
{"a": "int", "b": "int", "d": "int", "e": "int"}, df
)
assert str(e.value).startswith("Missing columns:")