Source code for message_ix_models.tests.model.water.data.test_irrigation

from message_ix import Scenario

from message_ix_models import ScenarioInfo
from message_ix_models.model.structure import get_codes
from message_ix_models.model.water.data.irrigation import add_irr_structure


[docs] def test_add_irr_structure(test_context): # FIXME You probably want this to be part of a common setup rather than writing # something like this for every test test_context.type_reg = "country" test_context.regions = "ZMB" nodes = get_codes(f"node/{test_context.regions}") nodes = list(map(str, nodes[nodes.index("World")].child)) test_context.map_ISO_c = {test_context.regions: nodes[0]} mp = test_context.get_platform() scenario_info = { "mp": mp, "model": "test water model", "scenario": "test water scenario", "version": "new", } s = Scenario(**scenario_info) s.add_horizon(year=[2020, 2030, 2040]) s.add_set("technology", ["tech1", "tech2"]) s.add_set("node", ["loc1", "loc2"]) s.add_set("year", [2020, 2030, 2040]) # FIXME same as above test_context["water build info"] = ScenarioInfo(s) # Call the function to be tested result = add_irr_structure(test_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", ] )