Skip to content

Commit b50fd0e

Browse files
committed
Add failing test for heterogenous input periods with divide
1 parent 1738912 commit b50fd0e

File tree

2 files changed

+27
-3
lines changed

2 files changed

+27
-3
lines changed

openfisca_core/holders.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,7 @@ def set_input_divide_by_period(holder, period, array):
330330
sub_period = period.start.period(cached_period_unit)
331331
while sub_period.start < after_instant:
332332
if holder.get_array(sub_period) is None:
333+
print("*** set", holder.variable.name, sub_period, divided_array)
333334
holder._set(sub_period, divided_array)
334335
sub_period = sub_period.offset(1)
335336
elif not (remaining_array == 0).all():

tests/core/test_simulation_builder.py

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,13 @@
1111
from openfisca_core.tools import assert_near
1212
from openfisca_core.tools.test_runner import yaml
1313
from openfisca_core.entities import Entity, GroupEntity
14+
from openfisca_core.holders import set_input_divide_by_period
1415
from openfisca_core.populations import Population
1516
from openfisca_core.variables import Variable
1617
from openfisca_country_template.entities import Household
1718
from openfisca_country_template.situation_examples import couple
1819
from openfisca_core.errors import SituationParsingError
19-
from openfisca_core.periods import ETERNITY
20+
from openfisca_core.periods import ETERNITY, MONTH
2021
from openfisca_core.indexed_enums import Enum as OFEnum
2122

2223

@@ -87,9 +88,16 @@ def __init__(self):
8788
@fixture
8889
def persons():
8990
class TestPersonEntity(Entity):
91+
def __init__(self, key, plural, label, doc):
92+
super().__init__(key, plural, label, doc)
93+
self.variables = {}
94+
9095
def get_variable(self, variable_name):
91-
result = TestVariable(self)
92-
result.name = variable_name
96+
result = self.variables.get(variable_name)
97+
if result is None:
98+
result = TestVariable(self)
99+
result.name = variable_name
100+
self.variables[variable_name] = result
93101
return result
94102

95103
def check_variable_defined_for_entity(self, variable_name):
@@ -267,6 +275,21 @@ def test_finalize_person_entity(simulation_builder, persons):
267275
assert population.ids == ['Alicia', 'Javier']
268276

269277

278+
def test_heterogenous_periods(simulation_builder, persons):
279+
persons_json = {'Alicia': {'salary': {'2018': 6000}}, 'Javier': {'salary': {'2018': 6000, '2018-01': 500}}}
280+
simulation_builder.add_person_entity(persons, persons_json)
281+
population = Population(persons)
282+
283+
salary = population.entity.get_variable('salary')
284+
salary.definition_period = MONTH
285+
salary.set_input = set_input_divide_by_period
286+
287+
simulation_builder.finalize_variables_init(population)
288+
assert_near(population.get_holder('salary').get_array('2018-01'), [500, 500])
289+
assert population.count == 2
290+
assert population.ids == ['Alicia', 'Javier']
291+
292+
270293
def test_canonicalize_period_keys(simulation_builder, persons):
271294
persons_json = {'Alicia': {'salary': {'year:2018-01': 100}}}
272295
simulation_builder.add_person_entity(persons, persons_json)

0 commit comments

Comments
 (0)