diff --git a/message_ix_models/tests/model/water/test_cooling.py b/message_ix_models/tests/model/water/test_cooling.py
new file mode 100644
index 0000000000..4cc833a016
--- /dev/null
+++ b/message_ix_models/tests/model/water/test_cooling.py
@@ -0,0 +1,41 @@
+
+import pandas as pd
+from unittest.mock import MagicMock, patch
+
+from message_ix_models.model.water.data.irrigation import non_cooling_tec
+
+def test_non_cooling_tec():
+    # Mock the context
+    context = {
+        "water build info": {"Y": [2020, 2030, 2040]},
+        "type_reg": "country",
+        "regions": "test_region",
+        "map_ISO_c": {"test_region": "test_ISO"},
+        "get_scenario": MagicMock(return_value=MagicMock(par=MagicMock(return_value=pd.DataFrame({
+            "technology": ["tech1", "tech2"],
+            "node_loc": ["loc1", "loc2"],
+            "node_dest": ["dest1", "dest2"],
+            "year_vtg": ["2020", "2020"],
+            "year_act": ["2020", "2020"],
+        }))))
+    }
+
+    # Mock the DataFrame read from CSV
+    df = pd.DataFrame({
+        "technology_group": ["cooling", "non-cooling"],
+        "technology_name": ["cooling_tech1", "non_cooling_tech1"],
+        "water_supply_type": ["freshwater_supply", "freshwater_supply"],
+        "water_withdrawal_mid_m3_per_output": [1, 2],
+    })
+    
+    # Mock the function 'private_data_path' to return the mocked DataFrame
+    with patch('message_ix_models.util.private_data_path', return_value="path/to/file"), \
+         patch('pandas.read_csv', return_value=df):
+        
+        # Call the function to be tested
+        result = non_cooling_tec(context)
+        
+        # Assert the results
+        assert isinstance(result, dict)
+        assert "input" in result
+        assert all(col in result["input"].columns for col in ["technology", "value", "unit", "level", "commodity", "mode", "time", "time_origin", "node_origin", "node_loc", "year_vtg", "year_act"])
diff --git a/message_ix_models/tests/model/water/test_irrigation.py b/message_ix_models/tests/model/water/test_irrigation.py
new file mode 100644
index 0000000000..3bb5252531
--- /dev/null
+++ b/message_ix_models/tests/model/water/test_irrigation.py
@@ -0,0 +1,33 @@
+import pandas as pd
+from unittest.mock import patch
+
+from message_ix_models.model.water.data.irrigation import add_irr_structure
+
+def test_add_irr_structure():
+    # Mock the context
+    context = {
+        "water build info": {"Y": [2020, 2030, 2040]},
+        "type_reg": "country",
+        "regions": "test_region",
+        "map_ISO_c": {"test_region": "test_ISO"},
+    }
+    
+    # Mock the DataFrame read from CSV
+    df_node = pd.DataFrame({
+        "BCU_name": ["1", "2"],
+        "REGION": ["region1", "region2"]
+    })
+    
+    # Mock the function 'private_data_path' to return the mocked DataFrame
+    with patch('message_ix_models.util.private_data_path', return_value="path/to/file"), \
+         patch('pandas.read_csv', return_value=df_node):
+        
+        # Call the function to be tested
+        result = add_irr_structure(context)
+        
+        # Assert the results
+        assert isinstance(result, dict)
+        assert "input" in result
+        assert "output" in result
+        assert all(col in result["input"].columns for col in ["technology", "value", "unit", "level", "commodity", "mode", "time", "time_origin", "node_origin", "node_loc", "year_vtg", "year_act"])
+        assert all(col in result["output"].columns for col in ["technology", "value", "unit", "level", "commodity", "mode", "time", "time_dest", "node_loc", "node_dest", "year_vtg", "year_act"])
diff --git a/message_ix_models/tests/model/water/test_nexus.py b/message_ix_models/tests/model/water/test_nexus.py
deleted file mode 100644
index 9b9a5e9bd9..0000000000
--- a/message_ix_models/tests/model/water/test_nexus.py
+++ /dev/null
@@ -1,169 +0,0 @@
-import pandas as pd
-import pytest
-from message_ix_models.model import bare
-
-from message_data.model.water import build
-from message_data.model.water.data.demands import (
-    add_irrigation_demand,
-    add_sectoral_demands,
-    add_water_availability,
-)
-from message_data.model.water.data.infrastructure import (
-    add_desalination,
-    add_infrastructure_techs,
-)
-from message_data.model.water.data.irrigation import add_irr_structure
-from message_data.model.water.data.water_for_ppl import cool_tech, non_cooling_tec
-from message_data.model.water.data.water_supply import add_e_flow, add_water_supply
-
-pytestmark = pytest.mark.xfail(reason="The tests will be updated in the preceeding PR")
-
-
-def configure_build(context, regions, years):
-    context.update(regions=regions, years=years)
-
-    # Information about the corresponding base model
-    info = bare.get_spec(context)["add"]
-    context["water build info"] = info
-    context["water spec"] = build.get_spec(context)
-
-    return info
-
-
-# def test_build(request, test_context):
-#     scenario = testing.bare_res(request, test_context)
-
-#     # Code runs on the bare RES
-#     build(scenario)
-
-#     # New set elements were added
-#     assert "extract_surfacewater" in scenario.set("technology").tolist()
-
-
-# def test_get_spec(session_context):
-#     # Code runs
-#     spec = get_spec()
-
-#     # Expected return type
-#     assert isinstance(spec, dict) and len(spec) == 3
-
-#     # Contents are read correctly
-#     assert "water_supply" in spec["require"].set["level"]
-
-
-# def test_read_config(session_context):
-#     # read_config() returns a reference to the current context
-#     context = read_config()
-#     assert context is session_context
-
-#     # 'add' elements have been converted to Code objects
-#     assert isinstance(context["water"]["set"]["technology"]["add"][0], Code)
-
-
-def test_read_data(test_context, regions, years):
-
-    context = test_context
-    # info = configure_build(context, regions, years)
-
-    data_irr = add_irr_structure(context)
-
-    # Returns a mapping
-    assert {"input", "output"} == set(data_irr.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_irr.values()))
-
-    data_water_ppl = cool_tech(context)
-
-    # Returns a mapping
-    assert {
-        "growth_activity_lo",
-        "growth_activity_up",
-        "input",
-        "output",
-        "capacity_factor",
-        "addon_conversion",
-        "addon_lo",
-        "inv_cost",
-        "historical_new_capacity",
-        "emission_factor",
-        "historical_activity",
-    } == set(data_water_ppl.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_water_ppl.values()))
-
-    data_non_cool = non_cooling_tec(context)
-
-    # Returns a mapping
-    assert {"input"} == set(data_non_cool.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_non_cool.values()))
-
-    data_infrastructure = add_infrastructure_techs(context)
-
-    # Returns a mapping
-    assert {
-        "input",
-        "output",
-        "capacity_factor",
-        "technical_lifetime",
-        "inv_cost",
-        "fix_cost",
-        "construction_time",
-    } == set(data_non_cool.keys())
-    assert all(
-        map(lambda df: isinstance(df, pd.DataFrame), data_infrastructure.values())
-    )
-
-    data_desal = add_desalination(context)
-    # Returns a mapping
-    assert {
-        "input",
-        "output",
-        "bound_total_capacity_up",
-        "historical_new_capacity",
-        "inv_cost",
-        "var_cost",
-        "bound_activity_lo",
-        "construction_time",
-        "bound_total_capacity_up",
-    } == set(data_desal.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_desal.values()))
-
-    data_demands = add_sectoral_demands(context)
-    # Returns a mapping
-    assert {"demand", "historical_new_capacity", "share_commodity_lo"} == set(
-        data_demands.keys()
-    )
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_demands.values()))
-
-    data_supply = add_water_supply(context)
-    # Returns a mapping
-    assert {
-        "input",
-        "output",
-        "historical_new_capacity",
-        "var_cost",
-        "share_mode_up",
-        "technical_lifetime",
-        "inv_cost",
-        "fix_cost",
-        "",
-    } == set(data_supply.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_supply.values()))
-
-    data_eflow = add_e_flow(context)
-    # Returns a mapping
-    assert {"bound_activity_lo"} == set(data_eflow.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_eflow.values()))
-
-    data_irr = add_irrigation_demand(context)
-    # Returns a mapping
-    assert {"land_input"} == set(data_irr.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_irr.values()))
-
-    data_irr = add_irrigation_demand(context)
-    # Returns a mapping
-    assert {"land_input"} == set(data_irr.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_irr.values()))
-
-    data_water_avail = add_water_availability(context)
-    # Returns a mapping
-    assert {"share_commodity_lo", "demand"} == set(data_water_avail.keys())
-    assert all(map(lambda df: isinstance(df, pd.DataFrame), data_water_avail.values()))
diff --git a/message_ix_models/tests/model/water/test_utils.py b/message_ix_models/tests/model/water/test_utils.py
new file mode 100644
index 0000000000..81ad7ae59f
--- /dev/null
+++ b/message_ix_models/tests/model/water/test_utils.py
@@ -0,0 +1,123 @@
+import pandas as pd
+import xarray as xr
+
+
+from message_ix_models.model.water import build, read_config, Context, map_add_on, add_commodity_and_level, map_yv_ya_lt
+from unittest.mock import patch
+
+
+def test_read_config():
+    # Mock the context
+    context = Context(0)
+
+    # Mock the data returned by load_private_data
+    mock_data = {"test_key": "test_value"}
+
+    # Mock the load_private_data function to return mock_data
+    with patch('message_ix_models.util.load_private_data', return_value=mock_data):
+        # Call the function to be tested
+        result = read_config(context)
+
+    # Assert the results
+    assert isinstance(result, Context)
+    assert result["water config"] == mock_data
+    assert result["water set"] == mock_data
+    assert result["water technology"] == mock_data
+
+
+def test_map_add_on():
+    # Mock the context
+    context = Context(0)
+
+    # Mock the data returned by read_config
+    mock_data = {
+        "water set": {
+            "add_on": {"add": [Code(id="1", name="test_1")]},
+            "type_addon": {"add": [Code(id="2", name="test_2")]}
+        }
+    }
+
+    # Mock the read_config function to return mock_data
+    with patch('your_module.read_config', return_value=mock_data):
+        # Call the function to be tested
+        result = map_add_on()
+
+    # Assert the results
+    expected = [Code(id="12", name="test_1, test_2")]
+    assert result == expected
+
+    # Testing with rtype = 'indexers'
+    with patch('your_module.read_config', return_value=mock_data):
+        result = map_add_on(rtype="indexers")
+
+    expected = {
+        'add_on': xr.DataArray(['1'], dims="consumer_group"),
+        'type_addon': xr.DataArray(['2'], dims="consumer_group"),
+        'consumer_group': xr.DataArray(['12'], dims="consumer_group"),
+    }
+    for key in expected:
+        assert (result[key] == expected[key]).all().item()
+
+
+
+def test_add_commodity_and_level():
+    # Mock the dataframe
+    df = pd.DataFrame({'technology': ['tech1', 'tech2']})
+
+    # Mock the data returned by Context.get_instance and get_codes
+    mock_context_data = {
+        "water set": {
+            "technology": {
+                "add": pd.Series(
+                    data=[Code(id="tech1", anno={"input": {"commodity": "com1", "level": "lev1"}}),
+                          Code(id="tech2", anno={"input": {"commodity": "com2"}})],
+                    name="tech"
+                )
+            }
+        }
+    }
+    mock_codes_data = pd.Series(
+        data=[Code(id="com1", anno={"level": "lev1"}),
+              Code(id="com2", anno={"level": "lev2"})],
+        name="com"
+    )
+
+    # Mock the Context.get_instance and get_codes functions to return mock_data
+    with patch('your_module.Context.get_instance', return_value=mock_context_data), \
+         patch('your_module.get_codes', return_value=mock_codes_data):
+        # Call the function to be tested
+        result = add_commodity_and_level(df)
+
+    # Assert the results
+    expected = pd.DataFrame({
+        'technology': ['tech1', 'tech2'],
+        'commodity': ['com1', 'com2'],
+        'level': ['lev1', 'lev2']
+    })
+    pd.testing.assert_frame_equal(result, expected)
+    
+    
+
+def test_map_yv_ya_lt():
+    periods = (2010, 2020, 2030, 2040)
+    lt = 20
+    ya = 2020
+
+    expected = pd.DataFrame({
+        'year_vtg': [2010, 2020, 2020, 2030],
+        'year_act': [2020, 2020, 2030, 2040]
+    })
+
+    result = map_yv_ya_lt(periods, lt, ya)
+
+    pd.testing.assert_frame_equal(result, expected)
+
+    # test with no active year specified
+    expected_no_ya = pd.DataFrame({
+        'year_vtg': [2010, 2020, 2020, 2030, 2030, 2040],
+        'year_act': [2020, 2020, 2030, 2030, 2040, 2040]
+    })
+
+    result_no_ya = map_yv_ya_lt(periods, lt)
+
+    pd.testing.assert_frame_equal(result_no_ya, expected_no_ya)
\ No newline at end of file
diff --git a/message_ix_models/tests/model/water/test_water_supply.py b/message_ix_models/tests/model/water/test_water_supply.py
new file mode 100644
index 0000000000..1de2955cea
--- /dev/null
+++ b/message_ix_models/tests/model/water/test_water_supply.py
@@ -0,0 +1,137 @@
+import pandas as pd
+from unittest.mock import patch
+from message_ix_models.model.water.data.water_supply import add_e_flow, add_water_supply, map_basin_region_wat
+
+
+
+def test_map_basin_region_wat():
+    # Mock the context
+    context = {
+        "water build info": {"Y": [2020, 2030, 2040]},
+        "type_reg": "country",
+        "regions": "test_region",
+        "map_ISO_c": {"test_region": "test_ISO"},
+        "RCP": "test_RCP",
+        "REL": "test_REL",
+        "time": "year",
+    }
+
+    # Mock the DataFrames read from CSV
+    df_x = pd.DataFrame({
+        "BCU_name": ["test_BCU"],
+    })
+    
+    df_sw = pd.DataFrame({
+        "Unnamed: 0": [0],
+        "BCU_name": ["test_BCU"],
+        "test_value": [1],
+    })
+
+    # Mock the function 'private_data_path' to return the mocked DataFrame
+    with patch('message_ix_models.util.private_data_path', return_value="path/to/file"), \
+         patch('pandas.read_csv', return_value=df_sw):
+        
+        # Call the function to be tested
+        result = map_basin_region_wat(context)
+        
+        # Assert the results
+        assert isinstance(result, pd.DataFrame)
+        assert all(col in result.columns for col in ["region", "mode", "date", "MSGREG", "share", "year", "time"])
+
+
+def test_add_water_supply():
+    # Mock the context
+    context = {
+        "water build info": {"Y": [2020, 2030, 2040]},
+        "type_reg": "country",
+        "regions": "test_region",
+        "map_ISO_c": {"test_region": "test_ISO"},
+        "RCP": "test_RCP",
+        "REL": "test_REL",
+        "time": "year",
+        "nexus_set": "nexus",
+        "get_scenario": lambda: {"firstmodelyear": 2020},
+    }
+
+    # Mock the DataFrames read from CSV
+    df_node = pd.DataFrame({
+        "BCU_name": ["test_BCU"],
+        "REGION": ["test_REGION"]
+    })
+
+    df_gwt = pd.DataFrame({
+        "REGION": ["test_REGION"],
+        "GW_per_km3_per_year": [1],
+    })
+
+    df_hist = pd.DataFrame({
+        "BCU_name": ["test_BCU"],
+        "hist_cap_sw_km3_year": [1],
+        "hist_cap_gw_km3_year": [1],
+    })
+
+    df_sw = pd.DataFrame({
+        "mode": ["test_mode"],
+        "MSGREG": ["test_MSGREG"],
+        "share": [1],
+        "year": [2020],
+        "time": ["year"],
+    })
+
+    # Mock the function 'private_data_path' to return the mocked DataFrame
+    with patch('message_ix_models.util.private_data_path', return_value="path/to/file"), \
+         patch('pandas.read_csv', return_value=df_node), \
+         patch('your_module.map_basin_region_wat', return_value=df_sw):
+
+        # Call the function to be tested
+        result = add_water_supply(context)
+
+        # Assert the results
+        assert isinstance(result, dict)
+        assert 'input' in result
+        assert 'output' in result
+        assert 'var_cost' in result
+        assert 'technical_lifetime' in result
+        assert 'inv_cost' in result
+
+        for df in result.values():
+            assert isinstance(df, pd.DataFrame)
+            
+            
+def test_add_e_flow():
+    # Mock the context
+    context = {
+        "water build info": {"Y": [2020, 2030, 2040]},
+        "regions": "test_region",
+        "RCP": "test_RCP",
+        "time": "year",
+        "SDG": True,
+    }
+
+    # Mock the DataFrames read from CSV
+    df_sw = pd.DataFrame({
+        "Region": ["test_Region"],
+        "value": [1],
+        "year": [2020],
+        "time": ["year"]
+    })
+
+    df_env = pd.DataFrame({
+        "Region": ["test_Region"],
+        "value": [1],
+        "year": [2020],
+        "time": ["year"]
+    })
+
+    # Mock the function 'read_water_availability' to return the mocked DataFrame
+    with patch('your_module.read_water_availability', return_value=(df_sw, df_sw)), \
+         patch('message_ix_models.util.private_data_path', return_value="path/to/file"), \
+         patch('pandas.read_csv', return_value=df_sw):
+
+        # Call the function to be tested
+        result = add_e_flow(context)
+
+        # Assert the results
+        assert isinstance(result, dict)
+        assert 'bound_activity_lo' in result
+        assert isinstance(result['bound_activity_lo'], pd.DataFrame)
\ No newline at end of file