def get_smoothed_spendable_income(dates, scheduled_income, allocations): """Push each scheduled income event as it occurs to an income smoother that smooths the spendable income (after allocations) from earlier income events to later income events. :param dates: list of datetime.datetimes, sequential in time :type dates: list of datetime.datetimes :param scheduled_income: list of scheduled income events :type scheduled_income: list of sub-classes of ScheduledEvent :param allocations: dict mapping (date, income) to list of expenses with metadata :type allocations: dict((datetime.datetime, sub-class of ScheduledEvent), [dicts with metadata]) :returns: dict mapping (date, income) to appropriate spendable amount (after allocations and savings for smoothing future income) :rtype: dict((datetime.datetime, sub-class of ScheduledEvent), float) """ start_date = dates[0] if dates else None last_date = dates[-1] if dates else None num_weeks = (last_date - start_date).days // 7 if dates else -1 weekly_nonallocated_income = [0.0] * (num_weeks + 1) for date in dates: for income in scheduled_income: if income.get_amount_on_date(date): week_index = (date - start_date).days // 7 weekly_nonallocated_income[week_index] += get_nonallocated_income( date, income, allocations ) smoother = IncomeSmoother() for week_index, nonallocated_income in enumerate(weekly_nonallocated_income): smoother.push(nonallocated_income, week_index) return smoother.to_dict()
class TestIncomeSmoother(unittest.TestCase): def setUp(self): self.balancer = IncomeSmoother() def tearDown(self): while not self.balancer.is_empty(): self.balancer.pop() def test_get_avg(self): self.balancer.push(2, SECONDARY_INDEX) assert self.balancer.get_running_avg() == 2.0 def test_push(self): self.balancer.push(5, SECONDARY_INDEX) (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 5 assert self.balancer.is_empty() def test_no_realloc(self): self.balancer.push(2, SECONDARY_INDEX) self.balancer.push(5, SECONDARY_INDEX) (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 5 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 2 assert self.balancer.is_empty() def test_realloc_simple(self): self.balancer.push(2, SECONDARY_INDEX) self.balancer.push(5, SECONDARY_INDEX) self.balancer.push(2, SECONDARY_INDEX) assert self.balancer.get_running_avg() == 3.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 4 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 3 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 2 assert self.balancer.is_empty() def test_realloc_complex(self): self.balancer.push(2, SECONDARY_INDEX) self.balancer.push(5, SECONDARY_INDEX) self.balancer.push(2, SECONDARY_INDEX) self.balancer.push(1, SECONDARY_INDEX) assert self.balancer.get_running_avg() == 2.5 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 3 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 2.5 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 2.5 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 2 assert self.balancer.is_empty() def test_realloc_all(self): self.balancer.push(100, SECONDARY_INDEX) self.balancer.push(100, SECONDARY_INDEX) self.balancer.push(100, SECONDARY_INDEX) self.balancer.push(100, SECONDARY_INDEX) self.balancer.push(0, SECONDARY_INDEX) assert self.balancer.get_running_avg() == 80.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 80.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 80.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 80.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 80.0 (spendable_amount, secondary_index) = self.balancer.pop() assert spendable_amount == 80.0 assert self.balancer.is_empty() def test_to_dict(self): self.balancer.push(2, SECONDARY_INDEX) self.balancer.push(5, SECONDARY_INDEX+1) tbl = self.balancer.to_dict() assert tbl[SECONDARY_INDEX] == 2 assert tbl[SECONDARY_INDEX+1] == 5 assert len(tbl) == 2