Add test_build

message_ix_models/model/water/build.py

@@ -174,7 +174,6 @@ def get_spec(context: Context) -> Mapping[str, ScenarioInfo]:
 
     return dict(require=require, remove=remove, add=add)
 
-    return dict(require=require, remove=remove, add=add)
 
 
 @lru_cache()

message_ix_models/tests/model/water/test_buid.py

@@ -0,0 +1,32 @@
+
+
+from message_ix_models import testing
+from message_ix_models.testing import bare_res
+
+from message_ix_models.model.water.build import build, get_spec,
+
+
+
+
+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"]
+
+
+

message_ix_models/tests/model/water/test_utils.py

@@ -6,117 +6,19 @@
 
 from message_ix_models import Context
 from message_ix_models.model.water.utils import (
-    add_commodity_and_level,
-    map_add_on,
     map_yv_ya_lt,
     read_config,
 )
 from message_ix_models.util import load_private_data
 
 
-def test_read_config(test_context):
-    # Mock the context
-    context = test_context
+def test_read_config(session_context):
+    # read_config() returns a reference to the current context
+    context = read_config()
+    assert context is session_context
 
-    # Call the function to be tested
-    result = read_config(context)
-
-    config_parts = ["water", "config.yaml"]
-    set_parts = ["water", "set.yaml"]
-    technology_parts = ["water", "technology.yaml"]
-
-    # Assert the results
-    assert isinstance(result, Context)
-    assert result["water config"] == load_private_data(*config_parts)
-    assert result["water set"] == load_private_data(*set_parts)
-    assert result["water technology"] == load_private_data(*technology_parts)
-
-
-def test_map_add_on():
-    # 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(
-        "message_ix_models.model.water.utils.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(
-        "message_ix_models.model.water.utils.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"]})
-
-    # FIXME Something here is seriously broken. Annotations need rework and
-    # please clarify what and how the annotations will be accessed and how the
-    # resulting data will be used!
-    # Mock the data returned by Context.get_instance and get_codes
-    mock_context_data = {
-        "water set": {
-            "technology": {
-                "add": [
-                    Code(
-                        id="tech1",
-                        annotations=[
-                            Annotation("input", "commodity", "com1", "level", "lev1")
-                        ],
-                    ),
-                    Code(
-                        id="tech2",
-                        annotations=[Annotation("input", "commodity", "com2")],
-                    ),
-                ],
-            }
-        }
-    }
-    mock_codes_data = [
-        Code(id="com1", annotations=[Annotation("level", "lev1")]),
-        Code(id="com2", annotations=[Annotation("level", "lev2")]),
-    ]
-
-    # Mock the Context.get_instance and get_codes functions to return mock_data
-    with patch(
-        "message_ix_models.util.context.Context.get_instance",
-        return_value=mock_context_data,
-    ), patch(
-        "message_ix_models.model.structure.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)
+    # 'add' elements have been converted to Code objects
+    assert isinstance(context["water"]["set"]["technology"]["add"][0], Code)
 
 
 def test_map_yv_ya_lt():