Add and test WWTW overrides

tests/test_wtw.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 from wsimod.core import constants
-from wsimod.nodes.wtw import WTW
+from wsimod.nodes.wtw import WTW, WWTW
 
 
 class MyTestClass(TestCase):
@@ -101,5 +101,18 @@ def test_override(self):
         wtw.apply_overrides(overrides)
         self.assertSetEqual(set(overrides.keys()), set(['name']))
         self.assertEqual(wtw.treatment_throughput_capacity, 20)
+
+    def test_wwtw_overrides(self):
+        wwtw = WWTW(name='')
+        vol = wwtw.process_parameters['volume']['constant']
+        wwtw.apply_overrides({'treatment_throughput_capacity' : 20,
+                              'process_parameters' : {'phosphate' : 
+                                                      {'constant' : 0.01}},
+                              'stormwater_storage_capacity': 100})
+        self.assertEqual(wwtw.treatment_throughput_capacity, 20)
+        self.assertEqual(wwtw.process_parameters['phosphate']['constant'], 0.01)
+        self.assertEqual(wwtw.process_parameters['volume']['constant'], vol)
+        self.assertEqual(wwtw.stormwater_storage_capacity, 100)
+
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/wtw.py

@@ -286,6 +286,23 @@ def __init__(
             lambda: self.ds_vqip(self.liquor, self.liquor_)
         )  # Change in liquor
 
+    def apply_overrides(self, overrides=Dict[str, Any]):
+        """Apply overrides to the stormwater tank and WWTW.
+        """
+        self.stormwater_storage_capacity = overrides.pop(
+            "stormwater_storage_capacity", 
+            self.stormwater_storage_capacity)
+        self.stormwater_storage_area = overrides.pop(
+            "stormwater_storage_area", 
+            self.stormwater_storage_area)
+        self.stormwater_storage_elevation = overrides.pop(
+            "stormwater_storage_elevation", 
+            self.stormwater_storage_elevation)
+        self.stormwater_tank.area = self.stormwater_storage_area
+        self.stormwater_tank.capacity = self.stormwater_storage_capacity
+        self.stormwater_tank.datum = self.stormwater_storage_elevation
+        super().apply_overrides(overrides)
+
     def calculate_discharge(self):
         """Clear stormwater tank if possible, and call treat_current_input."""
         # Run WWTW model