diff --git a/message_ix_models/tests/model/water/test_water_supply.py b/message_ix_models/tests/model/water/test_water_supply.py
index 134cd870c4..6d78f2cfa3 100644
--- a/message_ix_models/tests/model/water/test_water_supply.py
+++ b/message_ix_models/tests/model/water/test_water_supply.py
@@ -40,7 +40,7 @@ def test_map_basin_region_wat():
     with patch(
         "message_ix_models.util.private_data_path", return_value="path/to/file"
     ), patch("pandas.read_csv", return_value=df_sw):
-        context['time'] = 'year'
+        context["time"] = "year"
         result = map_basin_region_wat(context)
 
         # Assert the results
@@ -98,7 +98,7 @@ def test_add_water_supply():
         "message_ix_models.util.private_data_path", return_value="path/to/file"
     ), patch("pandas.read_csv", return_value=df_node), patch(
         "message_ix_models.model.water.water_supply.map_basin_region_wat",
-        return_value=df_sw  # Adjust this import
+        return_value=df_sw,  # Adjust this import
     ):
         # Call the function to be tested
         result = add_water_supply(context)
@@ -137,7 +137,7 @@ def test_add_e_flow():
     # Mock the function 'read_water_availability' to return the mocked DataFrame
     with patch(
         "message_ix_models.model.water.demands.read_water_availability",
-        return_value=(df_sw, df_sw)
+        return_value=(df_sw, df_sw),
     ), patch(
         "message_ix_models.util.private_data_path", return_value="path/to/file"
     ), patch(