def test_flag_nulls_raise(data_raw):
expander = DataFrameETL(check_null_cols='raise')
drop_cols = ['col1', 'col2']
# add a col of all nans
data_raw['nantastic'] = pd.Series([np.NaN] * 3)
with pytest.raises(RuntimeError):
expander._flag_nulls(data_raw, drop_cols)
def test_flag_nulls_warn(data_raw):
expander = DataFrameETL(check_null_cols='warn')
drop_cols = ['col1', 'col2']
drop_cols_2 = ['col1', 'col2', 'nantastic']
assert expander._flag_nulls(data_raw, []) == []
assert expander._flag_nulls(data_raw, drop_cols) == drop_cols
# add a col of all nans
data_raw['nantastic'] = pd.Series([np.NaN] * 3)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter('always')
# check that we don't add the col if it's already being dropped
assert expander._flag_nulls(data_raw, drop_cols_2) == drop_cols_2
assert len(w) == 0
assert expander._flag_nulls(data_raw, drop_cols) == drop_cols_2
assert len(w) == 1
assert issubclass(w[-1].category, UserWarning)
