def test_compute_cumsum_new_value_cols_error():
"""
It should compute cumsum with multiple value
"""
with pytest.raises(ParamsValueError) as exc_info:
compute_cumsum(
pd.DataFrame(),
id_cols=['City', 'Country', 'Region'],
reference_cols=['Date'],
value_cols=['population'],
new_value_cols=['a', 'b']
)
expected = "`value_cols` and `new_value_cols` needs to have the same number of elements"
assert str(exc_info.value) == expected
def test_compute_cumsum():
"""
It should compute cumsum
"""
input_df = pd.read_csv(os.path.join(fixtures_base_dir, 'compute_cumsum.csv'))
cumsum_df = compute_cumsum(
input_df,
id_cols=['City', 'Country', 'Region'],
reference_cols=['Date'],
value_cols=['population']
)
assert input_df['population_cumsum'].equals(cumsum_df['population'])
def prepare_data(data, mapping):
mapping = clean_mapping(mapping)
data = clean_data(data)
data = data.merge(mapping, on=BRAND, how='left')
# Add cumul by year
data_cum = compute_cumsum(
data,
id_cols=[YEAR, BRAND, BRAND_CLASS, PRODUCT_LINE, COUNTRY],
reference_cols=[DATE, MONTH],
value_cols=[VALUE])
data_cum['period_type'] = 'cumul'
data['period_type'] = 'current'
data = pd.concat([data_cum, data], sort=False).reset_index(drop=True)
return data
