Apply overrides in FWTW

tests/test_wtw.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 from wsimod.core import constants
-from wsimod.nodes.wtw import WTW, WWTW
+from wsimod.nodes.wtw import WTW, WWTW, FWTW
 
 
 class MyTestClass(TestCase):
@@ -128,6 +128,28 @@ def test_wwtw_overrides(self):
         self.assertEqual(wwtw.process_parameters["volume"]["constant"], vol)
         self.assertEqual(wwtw.stormwater_storage_capacity, 100)
 
+    def test_fwtw_overrides(self):
+        fwtw = FWTW(name="")
+        vol = fwtw.process_parameters["volume"]["constant"]
+        fwtw.apply_overrides(
+            {
+                "treatment_throughput_capacity": 20,
+                "process_parameters": {"phosphate": {"constant": 0.02}},
+                "service_reservoir_storage_capacity": 100,
+                "service_reservoir_storage_area": 34.7,
+                "service_reservoir_storage_elevation": 68.2,
+            }
+        )
+        self.assertEqual(fwtw.treatment_throughput_capacity, 20)
+        self.assertEqual(fwtw.process_parameters["phosphate"]["constant"], 0.02)
+        self.assertEqual(fwtw.process_parameters["volume"]["constant"], vol)
+        self.assertEqual(fwtw.service_reservoir_storage_capacity, 100)
+        self.assertEqual(fwtw.service_reservoir_tank.capacity, 100)
+        self.assertEqual(fwtw.service_reservoir_storage_area, 34.7)
+        self.assertEqual(fwtw.service_reservoir_tank.area, 34.7)
+        self.assertEqual(fwtw.service_reservoir_storage_elevation, 68.2)
+        self.assertEqual(fwtw.service_reservoir_tank.datum, 68.2)
+
 
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/wtw.py

@@ -543,6 +543,33 @@ def __init__(
         self.mass_balance_ds.append(lambda: self.service_reservoir_tank.ds())
         self.mass_balance_out.append(lambda: self.unpushed_sludge)
 
+    def apply_overrides(self, overrides=Dict[str, Any]):
+        """Apply overrides to the service reservoir tank and FWTW.
+
+        Enables a user to override any parameter of the service reservoir tank, and
+        then calls any overrides in WTW.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.service_reservoir_storage_capacity = overrides.pop(
+            "service_reservoir_storage_capacity",
+            self.service_reservoir_storage_capacity,
+        )
+        self.service_reservoir_storage_area = overrides.pop(
+            "service_reservoir_storage_area", self.service_reservoir_storage_area
+        )
+        self.service_reservoir_storage_elevation = overrides.pop(
+            "service_reservoir_storage_elevation",
+            self.service_reservoir_storage_elevation,
+        )
+
+        self.service_reservoir_tank.capacity = self.service_reservoir_storage_capacity
+        self.service_reservoir_tank.area = self.service_reservoir_storage_area
+        self.service_reservoir_tank.datum = self.service_reservoir_storage_elevation
+        super().apply_overrides(overrides)
+
     def treat_water(self):
         """Pulls water, aiming to fill service reservoirs, calls WTW
         treat_current_input, avoids deficit, sends liquor and solids to sewers."""