add tests (incomplete) for water

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"])

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"])

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)