Coverage for tests/test_agriculture_demand_projections.py: 100%
103 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-16 23:05 +0000
1"""Tests for agriculture demand projection helpers."""
3import csv
4import sys
5import types
6from pathlib import Path
8import pandas as pd
9from openpyxl import Workbook
11fake_stage_0_settings = types.ModuleType("prepare_times_nz.stage_0.stage_0_settings")
12fake_stage_0_settings.BASE_YEAR = 2023
13sys.modules.setdefault(
14 "prepare_times_nz.stage_0.stage_0_settings", fake_stage_0_settings
15)
17# pylint: disable=wrong-import-position
18from prepare_times_nz.stage_3.demand_projections import agriculture
20# pylint: disable= duplicate-code
21ASSUMPTION_COLUMNS = [
22 "SectorGroup",
23 "Sector",
24 "Scenario",
25 "Method",
26 "Workbook",
27 "SheetName",
28 "SourceCategory1",
29 "SourceCategory2",
30 "ConstantIndex",
31 "Note",
32]
34WORKBOOK_NAME = "mfe/Detailed-results-for-ERP2-projection-scenarios.xlsx"
35SECTOR_GROUP = "Agriculture, Forestry and Fishing"
37TEST_ASSUMPTIONS = [
38 {
39 "SectorGroup": SECTOR_GROUP,
40 "Sector": "Dairy Cattle Farming",
41 "Scenario": "Steady",
42 "Method": "Workbook",
43 "Workbook": WORKBOOK_NAME,
44 "SheetName": "Baseline high",
45 "SourceCategory1": "",
46 "SourceCategory2": "Total dairy cattle",
47 "ConstantIndex": "",
48 "Note": "",
49 },
50 {
51 "SectorGroup": SECTOR_GROUP,
52 "Sector": "Dairy Cattle Farming",
53 "Scenario": "Shift",
54 "Method": "Workbook",
55 "Workbook": WORKBOOK_NAME,
56 "SheetName": "Baseline low",
57 "SourceCategory1": "",
58 "SourceCategory2": "Total dairy cattle",
59 "ConstantIndex": "",
60 "Note": "",
61 },
62 {
63 "SectorGroup": SECTOR_GROUP,
64 "Sector": "Livestock Farming",
65 "Scenario": "Steady",
66 "Method": "Workbook",
67 "Workbook": WORKBOOK_NAME,
68 "SheetName": "Baseline",
69 "SourceCategory1": "",
70 "SourceCategory2": "Sheep and beef 'stock units'",
71 "ConstantIndex": "",
72 "Note": "",
73 },
74 {
75 "SectorGroup": SECTOR_GROUP,
76 "Sector": "Forestry and Logging",
77 "Scenario": "Steady",
78 "Method": "Workbook",
79 "Workbook": WORKBOOK_NAME,
80 "SheetName": "Baseline",
81 "SourceCategory1": "Forestry (million m3)",
82 "SourceCategory2": "Harvested timber (TRV)",
83 "ConstantIndex": "",
84 "Note": "",
85 },
86 {
87 "SectorGroup": SECTOR_GROUP,
88 "Sector": "Indoor Cropping",
89 "Scenario": "Steady",
90 "Method": "Workbook",
91 "Workbook": WORKBOOK_NAME,
92 "SheetName": "Baseline",
93 "SourceCategory1": "",
94 "SourceCategory2": "Horticulture",
95 "ConstantIndex": "",
96 "Note": "",
97 },
98 {
99 "SectorGroup": SECTOR_GROUP,
100 "Sector": "Fishing, Hunting and Trapping",
101 "Scenario": "Steady",
102 "Method": "Constant",
103 "Workbook": "",
104 "SheetName": "",
105 "SourceCategory1": "",
106 "SourceCategory2": "",
107 "ConstantIndex": 1,
108 "Note": "",
109 },
110]
113def make_test_erp_workbook(path: Path):
114 """Create a minimal ERP-like workbook with scenario sheets and year columns."""
115 wb = Workbook()
116 ws = wb.active
117 ws.title = "Baseline"
119 for sheet_name, dairy_2050, livestock_2050 in [
120 ("Baseline", 70, 80),
121 ("Baseline high", 85, 80),
122 ("Baseline low", 60, 65),
123 ]:
124 ws = (
125 wb[sheet_name]
126 if sheet_name in wb.sheetnames
127 else wb.create_sheet(sheet_name)
128 )
129 ws.append([None, None])
130 ws.append(["Scenario", None])
131 ws.append([sheet_name, None])
132 ws.append([None, None])
133 ws.append([None, None])
134 ws.append([None, None, 2023, 2025, 2030, 2050])
135 ws.append([None, "Total dairy cattle", 100, 90, 80, dairy_2050])
136 ws.append([None, "Sheep and beef 'stock units'", 100, 95, 90, livestock_2050])
137 ws.append(
138 ["Forestry (million m3)", "Harvested timber (TRV)", 100, 110, 130, 160]
139 )
140 ws.append([None, "Horticulture", 100, 105, 110, 120])
141 ws.append(["Other agriculture", "Total", 100, 102, 104, 106])
143 wb.save(path)
146def write_assumptions_csv(path: Path, rows):
147 """Write the agriculture assumptions test file from structured rows."""
148 with path.open("w", newline="", encoding="utf-8") as file_obj:
149 writer = csv.DictWriter(file_obj, fieldnames=ASSUMPTION_COLUMNS)
150 writer.writeheader()
151 writer.writerows(rows)
154def write_baseyear_demand_csv(path: Path):
155 """Write a minimal base-year demand file for agriculture projection tests."""
156 rows = [
157 {
158 "Sector": "Dairy Cattle Farming",
159 "CommodityOut": "MILK",
160 "Island": "NI",
161 "Technology": "DAIRY_TECH",
162 "EndUse": "Process",
163 "Variable": "InputEnergy",
164 "Unit": "PJ",
165 "Value": 10.0,
166 },
167 {
168 "Sector": "Dairy Cattle Farming",
169 "CommodityOut": "MILK",
170 "Island": "NI",
171 "Technology": "DAIRY_TECH",
172 "EndUse": "Process",
173 "Variable": "OutputEnergy",
174 "Unit": "PJ",
175 "Value": 5.0,
176 },
177 {
178 "Sector": "Fishing, Hunting and Trapping",
179 "CommodityOut": "FISH",
180 "Island": "SI",
181 "Technology": "FISH_TECH",
182 "EndUse": "Process",
183 "Variable": "InputEnergy",
184 "Unit": "PJ",
185 "Value": 3.0,
186 },
187 ]
188 path.parent.mkdir(parents=True, exist_ok=True)
189 pd.DataFrame(rows).to_csv(path, index=False)
192def setup_temp_agriculture_paths(tmp_path: Path, monkeypatch):
193 """Point the agriculture module at temporary stage data locations."""
194 stage_2_data = tmp_path / "stage_2_baseyear_data"
195 stage_3_data = tmp_path / "stage_3_scenario_data"
196 output = stage_3_data / "demand_projections"
198 monkeypatch.setattr(agriculture, "STAGE_2_DATA", stage_2_data)
199 monkeypatch.setattr(agriculture, "STAGE_3_DATA", stage_3_data)
200 monkeypatch.setattr(agriculture, "OUTPUT", output)
201 monkeypatch.setattr(agriculture, "OUTPUT_CHECKS", output / "checks")
203 return stage_2_data, stage_3_data
206def get_index(df, sector, scenario, year):
207 """Return one compiled index value for the requested sector/scenario/year."""
208 return df[
209 (df["Sector"] == sector) & (df["Scenario"] == scenario) & (df["Year"] == year)
210 ]["Index"].iloc[0]
213def test_get_agriculture_growth_indices_reads_workbook_mappings(tmp_path):
214 """Workbook-mapped and constant agriculture projections should compile correctly."""
215 external_data_dir = tmp_path / "external_data"
216 workbook_dir = external_data_dir / "mfe"
217 workbook_dir.mkdir(parents=True)
218 workbook_path = workbook_dir / "Detailed-results-for-ERP2-projection-scenarios.xlsx"
219 make_test_erp_workbook(workbook_path)
221 assumptions_path = tmp_path / "agriculture_demand_projections.csv"
222 write_assumptions_csv(assumptions_path, TEST_ASSUMPTIONS)
224 df = agriculture.get_agriculture_growth_indices(
225 assumptions_path=assumptions_path,
226 external_data_dir=external_data_dir,
227 )
229 assert get_index(df, "Dairy Cattle Farming", "Steady", 2025) == 0.9
230 assert get_index(df, "Dairy Cattle Farming", "Steady", 2024) == 0.95
231 assert get_index(df, "Dairy Cattle Farming", "Steady", 2050) == 0.85
232 assert get_index(df, "Dairy Cattle Farming", "Shift", 2050) == 0.6
233 assert get_index(df, "Livestock Farming", "Steady", 2050) == 0.8
234 assert get_index(df, "Indoor Cropping", "Steady", 2050) == 1.2
235 assert get_index(df, "Fishing, Hunting and Trapping", "Steady", 2042) == 1.0
238def test_get_energy_demand_projections_uses_temp_stage_data(tmp_path, monkeypatch):
239 """Energy projections should be computed from temp test data, not repo intermediates."""
240 stage_2_data, _ = setup_temp_agriculture_paths(tmp_path, monkeypatch)
242 external_data_dir = tmp_path / "external_data"
243 workbook_dir = external_data_dir / "mfe"
244 workbook_dir.mkdir(parents=True)
245 workbook_path = workbook_dir / "Detailed-results-for-ERP2-projection-scenarios.xlsx"
246 make_test_erp_workbook(workbook_path)
248 assumptions_path = tmp_path / "agriculture_demand_projections.csv"
249 write_assumptions_csv(assumptions_path, TEST_ASSUMPTIONS)
250 baseyear_path = (
251 stage_2_data / "ag_forest_fish" / "baseyear_ag_forest_fish_demand.csv"
252 )
254 write_baseyear_demand_csv(baseyear_path)
255 original_get_growth_indices = agriculture.get_agriculture_growth_indices
256 monkeypatch.setattr(
257 agriculture,
258 "get_agriculture_growth_indices",
259 lambda: original_get_growth_indices(
260 assumptions_path=assumptions_path,
261 external_data_dir=external_data_dir,
262 ),
263 )
265 df = agriculture.get_energy_demand_projections("InputEnergy")
267 dairy_2025 = df[
268 (df["Sector"] == "Dairy Cattle Farming")
269 & (df["Scenario"] == "Steady")
270 & (df["Year"] == 2025)
271 & (df["Variable"] == "InputEnergy")
272 ]["Value"].iloc[0]
273 fishing_2042 = df[
274 (df["Sector"] == "Fishing, Hunting and Trapping")
275 & (df["Scenario"] == "Steady")
276 & (df["Year"] == 2042)
277 & (df["Variable"] == "InputEnergy")
278 ]["Value"].iloc[0]
280 assert dairy_2025 == 9.0
281 assert fishing_2042 == 3.0
284def test_main_writes_outputs_to_temp_stage_data(tmp_path, monkeypatch):
285 """Main should save outputs under temp stage data rather than repo intermediates."""
286 stage_2_data, stage_3_data = setup_temp_agriculture_paths(tmp_path, monkeypatch)
288 external_data_dir = tmp_path / "external_data"
289 workbook_dir = external_data_dir / "mfe"
290 workbook_dir.mkdir(parents=True)
291 workbook_path = workbook_dir / "Detailed-results-for-ERP2-projection-scenarios.xlsx"
292 make_test_erp_workbook(workbook_path)
294 assumptions_path = tmp_path / "agriculture_demand_projections.csv"
295 write_assumptions_csv(assumptions_path, TEST_ASSUMPTIONS)
296 baseyear_path = (
297 stage_2_data / "ag_forest_fish" / "baseyear_ag_forest_fish_demand.csv"
298 )
300 write_baseyear_demand_csv(baseyear_path)
301 original_get_growth_indices = agriculture.get_agriculture_growth_indices
302 monkeypatch.setattr(
303 agriculture,
304 "get_agriculture_growth_indices",
305 lambda: original_get_growth_indices(
306 assumptions_path=assumptions_path,
307 external_data_dir=external_data_dir,
308 ),
309 )
311 agriculture.main()
313 assert (
314 stage_3_data / "demand_projections" / "agriculture_demand_index.csv"
315 ).exists()
316 assert (
317 stage_3_data / "demand_projections" / "checks" / "agriculture_input.csv"
318 ).exists()
319 assert (
320 stage_3_data / "demand_projections" / "checks" / "agriculture_output.csv"
321 ).exists()