Refactor with the make_variable fixture

Author: cescofran

tests/core/conftest.py

@@ -1,7 +1,5 @@
 # -*- coding: utf-8 -*-
 
-from enum import Enum
-from datetime import date
 from pytest import fixture
 
 from openfisca_country_template import CountryTaxBenefitSystem
@@ -12,65 +10,26 @@
 from openfisca_core.variables import Variable
 from openfisca_core.periods import ETERNITY
 from openfisca_core.entities import Entity, GroupEntity
-from openfisca_core.indexed_enums import Enum as OFEnum
 from openfisca_core.model_api import *  # noqa analysis:ignore
 
 
-class TestVariable(Variable):
-    definition_period = ETERNITY
-    value_type = float
-
-    def __init__(self, entity):
-        self.__class__.entity = entity
-        super().__init__()
-
-
-@fixture
-def tax_benefit_system():
-    return CountryTaxBenefitSystem()
-
-
-@fixture
-def simulation_builder():
-    return SimulationBuilder()
-
-
-@fixture
-def simulation(tax_benefit_system, simulation_builder):
-    return SimulationBuilder().build_default_simulation(tax_benefit_system)
-
+# PERIODS
 
 @fixture
 def period():
     return "2016-01"
 
 
-@fixture
-def year_period():
-    return "2016"
+# ENTITIES
 
 
-@fixture
-def make_simulation(tax_benefit_system, simulation_builder, period):
-    def _make_simulation(data):
-        simulation_builder.default_period = period
-        return simulation_builder.build_from_variables(tax_benefit_system, data)
-    return _make_simulation
-
-
-@fixture
-def make_isolated_simulation(simulation_builder, period):
-    def _make_simulation(tbs, data):
-        simulation_builder.default_period = period
-        return simulation_builder.build_from_variables(tbs, data)
-    return _make_simulation
-
+class TestVariable(Variable):
+    definition_period = ETERNITY
+    value_type = float
 
-@fixture
-def make_simulation_from_yaml(tax_benefit_system, simulation_builder):
-    def _make_simulation_from_yaml(simulation_yaml):
-        return simulation_builder.build_from_dict(tax_benefit_system, yaml.safe_load(simulation_yaml))
-    return _make_simulation_from_yaml
+    def __init__(self, entity):
+        self.__class__.entity = entity
+        super().__init__()
 
 
 @fixture
@@ -87,53 +46,6 @@ def check_variable_defined_for_entity(self, variable_name):
     return TestPersonEntity("person", "persons", "", "")
 
 
-@fixture
-def int_variable(persons):
-
-    class intvar(Variable):
-        definition_period = ETERNITY
-        value_type = int
-        entity = persons
-
-        def __init__(self):
-            super().__init__()
-
-    return intvar()
-
-
-@fixture
-def date_variable(persons):
-
-    class datevar(Variable):
-        definition_period = ETERNITY
-        value_type = date
-        entity = persons
-
-        def __init__(self):
-            super().__init__()
-
-    return datevar()
-
-
-@fixture
-def enum_variable():
-
-    class TestEnum(Variable):
-        definition_period = ETERNITY
-        value_type = OFEnum
-        dtype = 'O'
-        default_value = '0'
-        is_neutralized = False
-        set_input = None
-        possible_values = Enum('foo', 'bar')
-        name = "enum"
-
-        def __init__(self):
-            pass
-
-    return TestEnum()
-
-
 @fixture
 def group_entity():
     class Household(GroupEntity):
@@ -218,132 +130,68 @@ def couple():
             }
         }
 
-
-@fixture
-def simulation_single(simulation_builder, tax_benefit_system, single):
-    return simulation_builder.build_from_entities(tax_benefit_system, single)
+# TAX BENEFIt SYSTEMS
 
 
 @fixture
-def simulation_couple(simulation_builder, tax_benefit_system, couple):
-    return simulation_builder.build_from_entities(tax_benefit_system, couple)
-
-
-class simple_variable(Variable):
-    entity = Person
-    definition_period = MONTH
-    value_type = int
-
-
-class variable_with_calculate_output_add(Variable):
-    entity = Person
-    definition_period = MONTH
-    value_type = int
-    calculate_output = calculate_output_add
+def tax_benefit_system():
+    return CountryTaxBenefitSystem()
 
 
-class variable_with_calculate_output_divide(Variable):
-    entity = Person
-    definition_period = YEAR
-    value_type = int
-    calculate_output = calculate_output_divide
+# SIMULATIONS
 
 
 @fixture
-def simulation_single_with_variables(simulation_builder, tax_benefit_system, single):
-    tax_benefit_system.add_variables(
-        simple_variable,
-        variable_with_calculate_output_add,
-        variable_with_calculate_output_divide
-        )
-    return simulation_builder.build_from_entities(tax_benefit_system, single)
-
-
-# 1 <--> 2 with same period
-class variable1(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
-
-    def formula(person, period):
-        return person('variable2', period)
-
-
-class variable2(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
-
-    def formula(person, period):
-        return person('variable1', period)
-
-
-# 3 <--> 4 with a period offset
-class variable3(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
-
-    def formula(person, period):
-        return person('variable4', period.last_month)
-
+def simulation_builder():
+    return SimulationBuilder()
 
-class variable4(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
 
-    def formula(person, period):
-        return person('variable3', period)
+@fixture
+def simulation(tax_benefit_system, simulation_builder):
+    return SimulationBuilder().build_default_simulation(tax_benefit_system)
 
 
-# 5 -f-> 6 with a period offset
-#   <---
-class variable5(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
+@fixture
+def make_simulation_from_entities(tax_benefit_system, simulation_builder):
+    def _make_simulation_from_entities(entities):
+        return simulation_builder.build_from_entities(tax_benefit_system, entities)
+    return _make_simulation_from_entities
 
-    def formula(person, period):
-        variable6 = person('variable6', period.last_month)
-        return 5 + variable6
 
+@fixture
+def make_simulation_from_tbs_and_variables(simulation_builder, period):
+    def _make_simulation_from_variables_and_tbs(tbs, variables):
+        simulation_builder.default_period = period
+        return simulation_builder.build_from_variables(tbs, variables)
+    return _make_simulation_from_variables_and_tbs
 
-class variable6(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
 
-    def formula(person, period):
-        variable5 = person('variable5', period)
-        return 6 + variable5
+@fixture
+def make_simulation_from_yaml(tax_benefit_system, simulation_builder):
+    def _make_simulation_from_yaml(simulation_yaml):
+        return simulation_builder.build_from_dict(tax_benefit_system, yaml.safe_load(simulation_yaml))
+    return _make_simulation_from_yaml
 
 
-class variable7(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
+@fixture
+def simulation_single(make_simulation_from_entities, single):
+    return make_simulation_from_entities(single)
 
-    def formula(person, period):
-        variable5 = person('variable5', period)
-        return 7 + variable5
 
+@fixture
+def simulation_couple(make_simulation_from_entities, couple):
+    return make_simulation_from_entities(couple)
 
-# december cotisation depending on november value
-class cotisation(Variable):
-    value_type = int
-    entity = Person
-    definition_period = MONTH
 
-    def formula(person, period):
-        if period.start.month == 12:
-            return 2 * person('cotisation', period.last_month)
-        else:
-            return person.empty_array() + 1
+# VARIABLES
 
 
 @fixture
-def simulation_with_variables(tax_benefit_system, simulation_builder):
-    tax_benefit_system.add_variables(variable1, variable2, variable3, variable4,
-        variable5, variable6, variable7, cotisation)
-    return SimulationBuilder().build_default_simulation(tax_benefit_system)
+def make_variable():
+    def _make_variable(name, **new_methods_and_attrs):
+        methods_and_attrs = {}
+        default = dict(value_type = int, entity = Person, definition_period = MONTH)
+        methods_and_attrs.update(default)
+        methods_and_attrs.update(new_methods_and_attrs)
+        return type(name, (Variable, ), methods_and_attrs)
+    return _make_variable

tests/core/test_calculate_output.py

@@ -1,6 +1,33 @@
-from pytest import raises
+from pytest import fixture, raises
 
 from openfisca_core.tools import assert_near
+from openfisca_core.simulations import calculate_output_add
+from openfisca_core.simulations import calculate_output_divide
+from openfisca_core.periods import YEAR
+
+
+@fixture
+def simple_variable(make_variable):
+    return make_variable('simple_variable')
+
+
+@fixture
+def variable_with_calculate_output_add(make_variable):
+    return make_variable('variable_with_calculate_output_add', calculate_output = calculate_output_add)
+
+
+@fixture
+def variable_with_calculate_output_divide(make_variable):
+    return make_variable('variable_with_calculate_output_divide',
+            definition_period = YEAR,
+            calculate_output = calculate_output_divide)
+
+
+@fixture
+def simulation_single_with_variables(simulation_builder, tax_benefit_system, single,
+        simple_variable, variable_with_calculate_output_add, variable_with_calculate_output_divide):
+    tax_benefit_system.add_variables(simple_variable, variable_with_calculate_output_add, variable_with_calculate_output_divide)
+    return simulation_builder.build_from_entities(tax_benefit_system, single)
 
 
 def test_calculate_output_default(simulation_single_with_variables):

tests/core/test_countries.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from pytest import raises
+from pytest import fixture, raises
 
 from openfisca_core.variables import Variable
 from openfisca_core.periods import MONTH
@@ -10,6 +10,14 @@
 from openfisca_core.tools import assert_near
 
 
+@fixture
+def make_simulation(tax_benefit_system, simulation_builder, period):
+    def _make_simulation(variables):
+        simulation_builder.default_period = period
+        return simulation_builder.build_from_variables(tax_benefit_system, variables)
+    return _make_simulation
+
+
 def test_input_variable(make_simulation, period):
     simulation = make_simulation({'salary': 2000})
     assert_near(simulation.calculate('salary', period), [2000], absolute_error_margin = 0.01)
@@ -76,8 +84,8 @@ def test_input_with_wrong_period(make_simulation):
         make_simulation({'basic_income': {2015: 12000}})
 
 
-def test_variable_with_reference(tax_benefit_system, make_isolated_simulation):
-    simulation_base = make_isolated_simulation(tax_benefit_system, {'salary': 4000})
+def test_variable_with_reference(tax_benefit_system, make_simulation_from_tbs_and_variables):
+    simulation_base = make_simulation_from_tbs_and_variables(tax_benefit_system, {'salary': 4000})
 
     revenu_disponible_avant_reforme = simulation_base.calculate('disposable_income', "2016-01")
     assert(revenu_disponible_avant_reforme > 0)
@@ -90,7 +98,7 @@ def formula(household, period):
 
     tax_benefit_system.update_variable(disposable_income)
 
-    simulation_reform = make_isolated_simulation(tax_benefit_system, {'salary': 4000})
+    simulation_reform = make_simulation_from_tbs_and_variables(tax_benefit_system, {'salary': 4000})
     revenu_disponible_apres_reforme = simulation_reform.calculate('disposable_income', "2016-01")
 
     assert(revenu_disponible_apres_reforme == 0)