diff --git a/flee/InputGeography.py b/flee/InputGeography.py
index 71e2dd31..71824251 100644
--- a/flee/InputGeography.py
+++ b/flee/InputGeography.py
@@ -28,6 +28,7 @@ def __init__(self):
         self.conflicts = {}
         self.attributes = {}
 
+
     @check_args_type
     def ReadConflictInputCSV(self, csv_name: str) -> None:
         """
@@ -390,9 +391,9 @@ def StoreInputGeographyInEcosystem(self, e):
 
 
     @check_args_type
-    def UpdateAttributeZones(self, e, attribute_name: str, time: int, Debug: bool = False) -> None:
+    def UpdateLocationAttributes(self, e, attribute_name: str, time: int, Debug: bool = False) -> None:
         attrlist = self.attributes[attribute_name]
-        for i in range(0, e.locations):
+        for i in range(0, len(e.locations)):
             loc_name = e.locations[i].name
             if loc_name in attrlist:
                 e.locations[i].attributes[attribute_name] = attrlist[loc_name][time]
@@ -411,6 +412,10 @@ def AddNewConflictZones(self, e, time: int, Debug: bool = False) -> None:
             time (int): Description
             Debug (bool, optional): Description
         """
+
+        if SimulationSettings.move_rules["FloodRulesEnabled"] is True:
+            self.UpdateLocationAttributes(e, "flood_level", time)
+
         if len(SimulationSettings.ConflictInputFile) == 0:
             for loc in self.locations:
                 if "conflict" in loc[4].lower() and int(loc[5]) == time:
diff --git a/tests/test_dflee.py b/tests/test_dflee.py
new file mode 100644
index 00000000..9b8fa33d
--- /dev/null
+++ b/tests/test_dflee.py
@@ -0,0 +1,90 @@
+import os
+from flee import flee
+from flee import InputGeography
+
+
+"""
+Test cases for DFlee related functionalities.
+"""
+
+def test_read_flood_csv():
+
+    flee.SimulationSettings.ReadFromYML("empty.yml")
+
+    e = flee.Ecosystem()
+
+    ig = InputGeography.InputGeography()
+
+    ig.ReadLocationsFromCSV(csv_name=os.path.join("test_data", "test_input_csv/locations.csv"))
+
+    ig.ReadLinksFromCSV(csv_name=os.path.join("test_data", "test_input_csv/routes.csv"))
+
+    ig.ReadClosuresFromCSV(csv_name=os.path.join("test_data", "test_input_csv/closures.csv"))
+
+    ig.ReadAttributeInputCSV("flood_level", "int", os.path.join("test_data","test_input_csv","flood_level.csv"))
+
+    print(ig.attributes["flood_level"])
+
+    #ig.attributes["flood_level"] = {'A': [0, 0, 1, 1, 2, 1, 1, 1, 1, 1, 1], 'B': [1, 1, 1, 3, 1, 1, 0, 0, 0, 0, 1]}
+    assert ig.attributes["flood_level"]['A'][10] == 1
+    assert ig.attributes["flood_level"]['A'][4] == 2
+    assert ig.attributes["flood_level"]['B'][10] == 1
+    assert ig.attributes["flood_level"]['B'][0] == 1
+    assert ig.attributes["flood_level"]['B'][9] == 0
+
+
+    #e, lm = ig.StoreInputGeographyInEcosystem(e=e)
+
+
+
+
+def test_flood_level_location_attribute():
+    
+    flee.SimulationSettings.ReadFromYML("empty.yml")
+    flee.SimulationSettings.move_rules["FloodRulesEnabled"] = True
+
+    e = flee.Ecosystem()
+
+    ig = InputGeography.InputGeography()
+
+    ig.ReadLocationsFromCSV(csv_name=os.path.join("test_data", "test_input_csv/locations.csv"))
+
+    ig.ReadLinksFromCSV(csv_name=os.path.join("test_data", "test_input_csv/routes.csv"))
+
+    ig.ReadClosuresFromCSV(csv_name=os.path.join("test_data", "test_input_csv/closures.csv"))
+
+    ig.ReadAttributeInputCSV("flood_level", "int", os.path.join("test_data","test_input_csv","flood_level.csv"))
+
+    e, lm = ig.StoreInputGeographyInEcosystem(e=e)
+
+    end_time = 11
+
+    new_refs = 1
+
+    # Insert refugee agents
+    for _ in range(0, new_refs):
+        e.addAgent(location=lm["A"], attributes={})
+
+    for t in range(0, end_time):
+
+        ig.AddNewConflictZones(e, t)
+        if t == 0:
+            assert lm["A"].attributes["flood_level"] == 0
+            assert e.locations[1].attributes["flood_level"] == 1
+
+        # Propagate the model by one time step.
+        e.evolve()
+
+        print(lm["A"].attributes["flood_level"])
+        print(e.locations[1].attributes["flood_level"])
+
+        if t == 3:
+            assert lm["A"].attributes["flood_level"] == 1
+            assert e.locations[1].attributes["flood_level"] == 3
+    
+    assert lm["A"].attributes["flood_level"] == 1
+    assert e.locations[1].attributes["flood_level"] == 1
+
+
+if __name__ == "__main__":
+    pass