test(simulation): add test to calculate add

openfisca_core/variables/variable.py

@@ -14,6 +14,7 @@
 from openfisca_core import periods, tools
 from openfisca_core.entities import Entity, GroupEntity
 from openfisca_core.indexed_enums import Enum, EnumArray
+from openfisca_core.periods import DateUnit, Period
 
 from . import config, helpers
 
@@ -135,7 +136,7 @@ def __init__(self, baseline_variable=None):
             )
         self.entity = self.set(attr, "entity", required=True, setter=self.set_entity)
         self.definition_period = self.set(
-            attr, "definition_period", required=True, allowed_values=periods.DateUnit
+            attr, "definition_period", required=True, allowed_values=DateUnit
         )
         self.label = self.set(attr, "label", allowed_type=str, setter=self.set_label)
         self.end = self.set(attr, "end", allowed_type=str, setter=self.set_end)
@@ -373,7 +374,7 @@ def get_introspection_data(cls):
 
     def get_formula(
         self,
-        period: Union[Instant, periods.Period, str, int] = None,
+        period: Union[Instant, Period, str, int] = None,
     ) -> Optional[Formula]:
         """Returns the formula to compute the variable at the given period.
 
@@ -398,7 +399,7 @@ def get_formula(
                 1
             ]  # peekitem gets the 1st key-value tuple (the oldest start_date and formula). Return the formula.
 
-        if isinstance(period, periods.Period):
+        if isinstance(period, Period):
             instant = period.start
         else:
             try:

tests/core/test_calculate_output.py

@@ -2,27 +2,28 @@
 
 from openfisca_country_template import entities, situation_examples
 
-from openfisca_core import periods, simulations, tools
+from openfisca_core import simulations, tools
+from openfisca_core.periods import DateUnit
 from openfisca_core.simulations import SimulationBuilder
 from openfisca_core.variables import Variable
 
 
 class simple_variable(Variable):
     entity = entities.Person
-    definition_period = periods.MONTH
+    definition_period = DateUnit.MONTH
     value_type = int
 
 
 class variable_with_calculate_output_add(Variable):
     entity = entities.Person
-    definition_period = periods.MONTH
+    definition_period = DateUnit.MONTH
     value_type = int
     calculate_output = simulations.calculate_output_add
 
 
 class variable_with_calculate_output_divide(Variable):
     entity = entities.Person
-    definition_period = periods.YEAR
+    definition_period = DateUnit.YEAR
     value_type = int
     calculate_output = simulations.calculate_output_divide
 

tests/core/test_countries.py

@@ -2,6 +2,7 @@
 
 from openfisca_core import periods, populations, tools
 from openfisca_core.errors import VariableNameConflictError, VariableNotFoundError
+from openfisca_core.periods import DateUnit
 from openfisca_core.simulations import SimulationBuilder
 from openfisca_core.variables import Variable
 
@@ -100,7 +101,7 @@ def test_variable_with_reference(make_simulation, isolated_tax_benefit_system):
     assert result > 0
 
     class disposable_income(Variable):
-        definition_period = periods.MONTH
+        definition_period = DateUnit.MONTH
 
         def formula(household, period):
             return household.empty_array()
@@ -116,7 +117,7 @@ def formula(household, period):
 def test_variable_name_conflict(tax_benefit_system):
     class disposable_income(Variable):
         reference = "disposable_income"
-        definition_period = periods.MONTH
+        definition_period = DateUnit.MONTH
 
         def formula(household, period):
             return household.empty_array()