Allow passing a SectorModel instance to SectorModelBuilder

nismod · Jun 12, 2017 · 556985b · 556985b
1 parent 2fce960
commit 556985b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/smif/sector_model.py b/smif/sector_model.py
@@ -224,9 +224,9 @@ class SectorModelBuilder(object):
 
     """
 
-    def __init__(self, name):
+    def __init__(self, name, sector_model=None):
         self._sector_model_name = name
-        self._sector_model = None
+        self._sector_model = sector_model
         self.logger = logging.getLogger(__name__)
 
     def load_model(self, model_path, classname):

diff --git a/tests/test_sos_model.py b/tests/test_sos_model.py
@@ -56,8 +56,7 @@ def get_sos_model_only_scenario_dependencies(setup_region_data):
     _ws = WaterSupplySectorModel()
     _ws.name = 'water_supply'
 
-    ws_builder = SectorModelBuilder('water_supply')
-    ws_builder._sector_model = _ws
+    ws_builder = SectorModelBuilder('water_supply', _ws)
 
     ws_builder.create_initial_system([])
     ws_builder.add_regions(builder.sos_model.regions)
@@ -103,8 +102,7 @@ def get_sos_model_only_scenario_dependencies(setup_region_data):
 
     _ws2 = WaterSupplySectorModel()
     _ws2.name = 'water_supply_2'
-    ws2_builder = SectorModelBuilder('water_supply_2')
-    ws2_builder._sector_model = _ws2
+    ws2_builder = SectorModelBuilder('water_supply_2', _ws2)
 
     ws2_builder.add_inputs([
         {