diff --git a/tests/test_wtw.py b/tests/test_wtw.py index 447145e..44fe486 100644 --- a/tests/test_wtw.py +++ b/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() diff --git a/wsimod/nodes/wtw.py b/wsimod/nodes/wtw.py index afd6c6a..4a957c2 100644 --- a/wsimod/nodes/wtw.py +++ b/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."""