diff --git a/hydromt_fiat/fiat.py b/hydromt_fiat/fiat.py
index 75d09603..b834729d 100644
--- a/hydromt_fiat/fiat.py
+++ b/hydromt_fiat/fiat.py
@@ -253,7 +253,7 @@ def setup_vulnerability_from_csv(self, csv_fn: Union[str, Path], unit: str) -> N
 
     def setup_road_vulnerability(
         self,
-        unit: str,
+        vertical_unit: str,
         threshold_value: float = 0.6,
         min_hazard_value: float = 0,
         max_hazard_value: float = 10,
@@ -261,7 +261,7 @@ def setup_road_vulnerability(
     ):
         if not self.vulnerability:
             self.vulnerability = Vulnerability(
-                unit,
+                vertical_unit,
                 self.logger,
             )
         self.vulnerability.create_step_function(
diff --git a/hydromt_fiat/workflows/exposure_vector.py b/hydromt_fiat/workflows/exposure_vector.py
index ef442012..f5fa3010 100644
--- a/hydromt_fiat/workflows/exposure_vector.py
+++ b/hydromt_fiat/workflows/exposure_vector.py
@@ -221,16 +221,16 @@ def setup_roads(
             # add the function to segmentize the roads into certain segments
 
         # Add the Primary Object Type and damage function, which is currently not set up to be flexible
-        roads["Primary Object Type"] = "road"
-        roads["Damage Function: Structure"] = "road"
+        roads["Primary Object Type"] = "roads"
+        roads["Damage Function: Structure"] = "roads"
 
         self.logger.info(
-            "The damage function 'road' is selected for all of the structure damage to the roads."
+            "The damage function 'roads' is selected for all of the structure damage to the roads."
         )
 
-        # Add the max potential damage to the roads
+        # Add the max potential damage and the length of the segments to the roads
         road_damage = self.data_catalog.get_dataframe(road_damage)
-        get_max_potential_damage_roads(roads, road_damage)
+        roads[["Max Potential Damage: Structure", "Segment Length [m]"]] = get_max_potential_damage_roads(roads, road_damage)
 
         self.set_exposure_geoms(roads[["Object ID", "geometry"]])
         self.set_geom_names("roads")
diff --git a/hydromt_fiat/workflows/roads.py b/hydromt_fiat/workflows/roads.py
index 9f2d84cd..a7924f14 100644
--- a/hydromt_fiat/workflows/roads.py
+++ b/hydromt_fiat/workflows/roads.py
@@ -39,4 +39,4 @@ def get_max_potential_damage_roads(
             "You are using the wrong unit for the segment length. Please use <'foot/feet/ft'> or <'meter/metre/m'>"
         )
 
-    return roads
+    return roads[["maximum_potential_damage", "segment_length"]]
diff --git a/hydromt_fiat/workflows/vulnerability.py b/hydromt_fiat/workflows/vulnerability.py
index 6b3e1693..d36bf891 100644
--- a/hydromt_fiat/workflows/vulnerability.py
+++ b/hydromt_fiat/workflows/vulnerability.py
@@ -308,6 +308,8 @@ def create_step_function(
 
         self.add(name, hazard_values, fraction_values)
 
+        # Set the area extraction method
+        self.set_area_extraction_methods()
 
     def get_damage_function_names(self):
         return list(self.functions.keys())
diff --git a/tests/test_roads.py b/tests/test_roads.py
index 04d55b3d..8d8d48c1 100644
--- a/tests/test_roads.py
+++ b/tests/test_roads.py
@@ -4,6 +4,7 @@
 import pytest
 import shutil
 import geopandas as gpd
+import pandas as pd
 
 
 EXAMPLEDIR = Path(
@@ -38,11 +39,11 @@
         "dir": "test_roads_from_OSM",
         "configuration": {
             "setup_road_vulnerability": {
-                "name" : "roads",
-                "theshold_value" : 0.5,
-                "min_hazard__value" : 0,
-                "max_hazard_value" :15,
-                "step_hazard_value" : 1,
+                "vertical_unit": "ft",
+                "threshold_value": 0.5,
+                "min_hazard_value": 0,
+                "max_hazard_value": 15,
+                "step_hazard_value": 1,
             },
             "setup_exposure_roads": {
                 "roads_fn": "OSM",
@@ -77,6 +78,13 @@ def test_setup_roads(case):
     assert root.joinpath("exposure", "exposure.csv").exists()
     assert root.joinpath("exposure", "region.gpkg").exists()
 
+    # Read the resulting exposure data and check if the required columns exist
+    exposure = pd.read_csv(root.joinpath("exposure", "exposure.csv"))
+    required_columns = ['Secondary Object Type', 'Object Name', 'lanes', 'Object ID', 
+                        'Primary Object Type', 'Damage Function: Structure',
+                        'Max Potential Damage: Structure', 'Segment Length [m]']
+    assert set(required_columns) == set(exposure.columns)
+
     # Check if the vulnerability data exists
     assert root.joinpath("vulnerability", "vulnerability_curves.csv").exists()