From b649e363001e5089f9483273d1aece82332bff65 Mon Sep 17 00:00:00 2001
From: liuly12 <leyang.liu@hotmail.com>
Date: Sat, 16 Mar 2024 14:22:44 +0000
Subject: [PATCH] Apply overrides in Distribution

---
 tests/test_distribution.py   |  4 ++++
 wsimod/nodes/distribution.py | 25 +++++++++++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index aeaabc7..239400e 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -67,6 +67,10 @@ def test_leakage(self):
         self.assertEqual(v2, arc1.vqip_in["volume"])
         self.assertEqual(v2 * 0.2, arc2.vqip_in["volume"])
 
+    def test_distribution_overrides(self):
+        distribution = Distribution(name="", leakage=0.2)
+        distribution.apply_overrides({'leakage': 0})
+        self.assertEqual(distribution.leakage, 0)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/wsimod/nodes/distribution.py b/wsimod/nodes/distribution.py
index 63a7cc0..fece94d 100644
--- a/wsimod/nodes/distribution.py
+++ b/wsimod/nodes/distribution.py
@@ -6,7 +6,7 @@
 
 from wsimod.core import constants
 from wsimod.nodes.nodes import Node
-
+from typing import Any, Dict
 
 def decorate_leakage_set(self, f):
     """Decorator to extend the functionality of `f` by introducing leakage. This is
@@ -125,15 +125,32 @@ def __init__(self, leakage=0, **kwargs):
         # Update handlers
         self.push_set_handler["default"] = self.push_set_deny
         self.push_check_handler["default"] = self.push_check_deny
-
-        if leakage > 0:
+        self.decorate_pull_handlers()
+    
+    def decorate_pull_handlers(self):
+        """Decorate handlers if there is leakage ratio.
+        """
+        if self.leakage > 0:
             self.pull_set_handler["default"] = decorate_leakage_set(
                 self, self.pull_set_handler["default"]
             )
             self.pull_check_handler["default"] = decorate_leakage_check(
                 self, self.pull_check_handler["default"]
             )
-
+    
+    def apply_overrides(self, overrides: Dict[str, Any] = {}):
+        """Apply overrides to the sewer.
+    
+        Enables a user to override any of the following parameters: 
+        leakage.
+        
+        Args:
+            overrides (dict, optional): Dictionary of overrides. Defaults to {}.
+        """
+        self.leakage = overrides.pop("leakage", 
+                                      self.leakage)
+        self.decorate_pull_handlers()
+        super().apply_overrides(overrides)
 
 class UnlimitedDistribution(Distribution):
     """"""