def test_add_axis_on_persons_distributes_roles(persons, households): simulation_builder = SimulationBuilder() simulation_builder.add_person_entity(persons, {'Alicia': {}, 'Javier': {}, 'Tom': {}}) simulation_builder.add_group_entity('persons', ['Alicia', 'Javier', 'Tom'], households, { 'housea': {'parents': ['Alicia']}, 'houseb': {'parents': ['Tom'], 'children': ['Javier']}, }) simulation_builder.register_variable('salary', persons) simulation_builder.add_parallel_axis({'count': 2, 'name': 'salary', 'min': 0, 'max': 3000, 'period': '2018-11'}) simulation_builder.expand_axes() assert [role.key for role in simulation_builder.get_roles('households')] == ['parent', 'child', 'parent', 'parent', 'child', 'parent']
def test_add_group_entity_loose_syntax(households): simulation_builder = SimulationBuilder() simulation_builder.add_group_entity( 'persons', ['Alicia', 'Javier', 'Sarah', '1'], households, { 'Household_1': { 'parents': ['Alicia', 'Javier'] }, 'Household_2': { 'parents': 1, 'children': 'Sarah' }, }) assert simulation_builder.get_count('households') == 2 assert simulation_builder.get_ids('households') == [ 'Household_1', 'Household_2' ] assert simulation_builder.get_memberships('households') == [0, 0, 1, 1] assert [role.key for role in simulation_builder.get_roles('households') ] == ['parent', 'parent', 'child', 'parent']