add test and fix

tests/test_extensions.py

@@ -201,6 +201,18 @@ def test_custom_class_from_file():
         model.run(dates=[to_datetime("2000-01-01")])
         assert model.nodes["node_name"].custom_attr == 2
 
+        model.save(temp_dir, "new_config.yml")
+
+        # Remove the custom class from the registry to test loading it again
+        del model
+        NODES_REGISTRY.pop("CustomNode", None)
+
+        model = Model()
+        model.load(temp_dir, "new_config.yml")
+        assert model.nodes["node_name"].custom_attr == 1
+        model.run(dates=[to_datetime("2000-01-01")])
+        assert model.nodes["node_name"].custom_attr == 2
+
 
 def test_custom_class_on_the_fly():
     """Test a custom class."""

wsimod/orchestration/model.py

@@ -140,6 +140,7 @@ def __init__(self):
         # self.arcs_type = {} #not sure that this would be necessary
         self.nodes = {}
         self.nodes_type = {}
+        self.extensions = []
 
         # Default orchestration
         self.orchestration = [
@@ -200,6 +201,7 @@ def load(self, address, config_name="config.yml", overrides={}):
             E.G. ADDITION FOR NEW ORCHESTRATION
         """
         load_extension_files(data.get("extensions", []))
+        self.extensions = data.get("extensions", [])
 
         if "orchestration" in data.keys():
             # Update orchestration
@@ -330,6 +332,7 @@ def save(self, address, config_name="config.yml", compress=False):
             "additive_pollutants": constants.ADDITIVE_POLLUTANTS,
             "non_additive_pollutants": constants.NON_ADDITIVE_POLLUTANTS,
             "float_accuracy": constants.FLOAT_ACCURACY,
+            "extensions": self.extensions,
         }
         if hasattr(self, "dates"):
             data["dates"] = [str(x) for x in self.dates]