From 8c74fa1fb5e661264019ff1aaeb7ec7e841573df Mon Sep 17 00:00:00 2001 From: Muhammad Awais Date: Sun, 7 Apr 2024 13:09:29 -0700 Subject: [PATCH] Add test_build --- message_ix_models/model/water/build.py | 1 - .../tests/model/water/test_buid.py | 32 +++++ .../tests/model/water/test_utils.py | 110 +----------------- 3 files changed, 38 insertions(+), 105 deletions(-) create mode 100644 message_ix_models/tests/model/water/test_buid.py diff --git a/message_ix_models/model/water/build.py b/message_ix_models/model/water/build.py index 7b2ffc6e23..b4becc9b68 100644 --- a/message_ix_models/model/water/build.py +++ b/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() diff --git a/message_ix_models/tests/model/water/test_buid.py b/message_ix_models/tests/model/water/test_buid.py new file mode 100644 index 0000000000..8c98f94caf --- /dev/null +++ b/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"] + + + diff --git a/message_ix_models/tests/model/water/test_utils.py b/message_ix_models/tests/model/water/test_utils.py index c754b1cdfb..9e04d5fc7e 100644 --- a/message_ix_models/tests/model/water/test_utils.py +++ b/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():