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)