Apply overrides in Demand and its subclasses

tests/test_demand.py

@@ -119,6 +119,48 @@ def test_garden_demand(self):
         d1 = {"volume": 0.03 * 0.4, "temperature": 0, "phosphate": 0}
         self.assertDictAlmostEqual(d1, reply)
 
+    def test_demand_overrides(self):
+        demand = Demand(
+            name="",
+            constant_demand=10,
+            pollutant_load={"phosphate": 0.1, "temperature": 12},
+        )
+        demand.apply_overrides(
+            {"constant_demand": 20, "pollutant_load": {"phosphate": 0.5}}
+        )
+        self.assertEqual(demand.constant_demand, 20)
+        self.assertDictEqual(
+            demand.pollutant_load, {"phosphate": 0.5, "temperature": 12}
+        )
+
+    def test_residentialdemand_overrides(self):
+        demand = ResidentialDemand(
+            name="",
+            gardening_efficiency=0.4,
+            pollutant_load={"phosphate": 0.1, "temperature": 12},
+        )
+        demand.apply_overrides(
+            {
+                "gardening_efficiency": 0.5,
+                "population": 153.2,
+                "per_capita": 32.4,
+                "constant_weighting": 47.5,
+                "constant_temp": 0.71,
+                "constant_demand": 20,
+                "pollutant_load": {"phosphate": 0.5},
+            }
+        )
+        self.assertEqual(demand.gardening_efficiency, 0.5)
+        self.assertEqual(demand.population, 153.2)
+        self.assertEqual(demand.per_capita, 32.4)
+        self.assertEqual(demand.constant_weighting, 47.5)
+        self.assertEqual(demand.constant_temp, 0.71)
+
+        self.assertEqual(demand.constant_demand, 20)
+        self.assertDictEqual(
+            demand.pollutant_load, {"phosphate": 0.5, "temperature": 12}
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/demand.py

@@ -5,6 +5,8 @@
 
 Converted to totals BD 2022-05-03
 """
+from typing import Any, Dict
+
 from wsimod.core import constants
 from wsimod.nodes.nodes import Node
 
@@ -28,7 +30,7 @@ def __init__(
                 is used. Defaults to 0.
             pollutant_load (dict, optional): Pollutant mass per timestep of
             constant_demand.
-                Defaults to {}.
+                Defaults to 0.
             data_input_dict (dict, optional):  Dictionary of data inputs relevant for
                 the node (temperature). Keys are tuples where first value is the name of
                 the variable to read from the dict and the second value is the time.
@@ -61,6 +63,19 @@ def __init__(
         self.mass_balance_out.append(lambda: self.total_backup)
         self.mass_balance_out.append(lambda: self.total_received)
 
+    def apply_overrides(self, overrides: Dict[str, Any] = {}):
+        """Apply overrides to the sewer.
+
+        Enables a user to override any of the following parameters:
+        constant_demand, pollutant_load.
+
+        Args:
+            overrides (dict, optional): Dictionary of overrides. Defaults to {}.
+        """
+        self.constant_demand = overrides.pop("constant_demand", self.constant_demand)
+        self.pollutant_load.update(overrides.pop("pollutant_load", {}))
+        super().apply_overrides(overrides)
+
     def create_demand(self):
         """Function to call get_demand, which should return a dict with keys that match
         the keys in directions.
@@ -198,6 +213,26 @@ def __init__(
         # Label as Demand class so that other nodes treat it the same
         self.__class__.__name__ = "Demand"
 
+    def apply_overrides(self, overrides: Dict[str, Any] = {}):
+        """Apply overrides to the sewer.
+
+        Enables a user to override any of the following parameters:
+        gardening_efficiency, population, per_capita, constant_weighting, constant_temp.
+
+        Args:
+            overrides (dict, optional): Dictionary of overrides. Defaults to {}.
+        """
+        self.gardening_efficiency = overrides.pop(
+            "gardening_efficiency", self.gardening_efficiency
+        )
+        self.population = overrides.pop("population", self.population)
+        self.per_capita = overrides.pop("per_capita", self.per_capita)
+        self.constant_weighting = overrides.pop(
+            "constant_weighting", self.constant_weighting
+        )
+        self.constant_temp = overrides.pop("constant_temp", self.constant_temp)
+        super().apply_overrides(overrides)
+
     def get_demand(self):
         """Overwrite get_demand and replace with custom functions.