From 477be1019a9246be6d91982a1c2d9e04bb2b7ec3 Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Thu, 24 Sep 2026 18:37:02 -0600 Subject: [PATCH] add MemoryMassBalanceOutputProvider --- src/itzi_core/providers/memory_output.py | 22 +++++++++++-- tests/test_mass_balance_provider.py | 42 ++++++------------------ tests/test_rainfall_precision.py | 14 ++------ 3 files changed, 33 insertions(+), 45 deletions(-) diff --git a/src/itzi_core/providers/memory_output.py b/src/itzi_core/providers/memory_output.py index 243f85b..1fb75ab 100644 --- a/src/itzi_core/providers/memory_output.py +++ b/src/itzi_core/providers/memory_output.py @@ -18,8 +18,16 @@ import numpy as np -from itzi_core.data_containers import DrainageNetworkAttributes, DrainageNetworkTopology -from itzi_core.providers.base import RasterOutputProvider, VectorOutputProvider +from itzi_core.data_containers import ( + DrainageNetworkAttributes, + DrainageNetworkTopology, + MassBalanceData, +) +from itzi_core.providers.base import ( + MassBalanceOutputProvider, + RasterOutputProvider, + VectorOutputProvider, +) class MemoryRasterOutputProvider(RasterOutputProvider): @@ -38,6 +46,16 @@ def write_arrays( ) +class MemoryMassBalanceOutputProvider(MassBalanceOutputProvider): + """Save mass-balance reports in memory.""" + + def __init__(self) -> None: + self.reports: list[MassBalanceData] = [] + + def log(self, report_data: MassBalanceData) -> None: + self.reports.append(report_data) + + class MemoryVectorOutputProvider(VectorOutputProvider): """Save drainage simulation outputs in memory.""" diff --git a/tests/test_mass_balance_provider.py b/tests/test_mass_balance_provider.py index bfde6cd..23f8574 100644 --- a/tests/test_mass_balance_provider.py +++ b/tests/test_mass_balance_provider.py @@ -16,35 +16,20 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING +from unittest.mock import Mock import numpy as np from itzi_core.const import TemporalType from itzi_core.data_containers import MassBalanceData, SimulationConfig, SurfaceFlowParameters from itzi_core.providers.base import MassBalanceOutputProvider +from itzi_core.providers.memory_output import MemoryMassBalanceOutputProvider from itzi_core.simulation_builder import SimulationBuilder if TYPE_CHECKING: from itzi_core.simulation import Simulation -class RecordingMassBalanceProvider(MassBalanceOutputProvider): - def __init__(self) -> None: - self.reports: list[MassBalanceData] = [] - self.finalize_calls = 0 - - def log(self, report_data: MassBalanceData) -> None: - self.reports.append(report_data) - - def finalize(self) -> None: - self.finalize_calls += 1 - - -class FalseyRecordingMassBalanceProvider(RecordingMassBalanceProvider): - def __bool__(self) -> bool: - return False - - def _build_simulation( domain_5by5, helpers, @@ -78,28 +63,19 @@ def _build_simulation( return simulation -def test_custom_provider_receives_simulation_mass_balance(domain_5by5, helpers) -> None: - provider = RecordingMassBalanceProvider() +def test_memory_provider_receives_simulation_mass_balance(domain_5by5, helpers) -> None: + provider = MemoryMassBalanceOutputProvider() simulation = _build_simulation(domain_5by5, helpers, provider=provider) simulation.initialize() simulation.update() + assert isinstance(provider, MassBalanceOutputProvider) assert len(provider.reports) == 2 assert all(isinstance(report, MassBalanceData) for report in provider.reports) assert provider.reports[-1].simulation_time == timedelta(seconds=1) -def test_falsey_injected_provider_is_used(domain_5by5, helpers) -> None: - provider = FalseyRecordingMassBalanceProvider() - - simulation = _build_simulation(domain_5by5, helpers, provider=provider) - simulation.initialize() - - assert simulation.report.mass_balance_output_provider is provider - assert len(provider.reports) == 1 - - def test_mass_balance_output_is_disabled_without_configuration(domain_5by5, helpers) -> None: simulation = _build_simulation(domain_5by5, helpers) @@ -108,11 +84,13 @@ def test_mass_balance_output_is_disabled_without_configuration(domain_5by5, help assert simulation.report.mass_balance_output_provider is None -def test_finalize_calls_mass_balance_provider_once(domain_5by5, helpers) -> None: - provider = FalseyRecordingMassBalanceProvider() +def test_finalize_calls_mass_balance_provider_once(domain_5by5, helpers, monkeypatch) -> None: + provider = MemoryMassBalanceOutputProvider() + finalize = Mock() + monkeypatch.setattr(provider, "finalize", finalize) simulation = _build_simulation(domain_5by5, helpers, provider=provider) simulation.initialize() simulation.finalize() - assert provider.finalize_calls == 1 + assert finalize.call_count == 1 diff --git a/tests/test_rainfall_precision.py b/tests/test_rainfall_precision.py index 8856a37..0cd8cb4 100644 --- a/tests/test_rainfall_precision.py +++ b/tests/test_rainfall_precision.py @@ -24,8 +24,8 @@ from itzi_core.compute import rastermetrics from itzi_core.const import InfiltrationModelType, TemporalType from itzi_core.data_containers import MassBalanceData, SimulationConfig, SurfaceFlowParameters -from itzi_core.providers.base import MassBalanceOutputProvider from itzi_core.providers.memory_output import ( + MemoryMassBalanceOutputProvider, MemoryRasterOutputProvider, MemoryVectorOutputProvider, ) @@ -34,14 +34,6 @@ RAIN_RATE = 10.0 / (1000 * 3600) -class CaptureMassBalanceOutputProvider(MassBalanceOutputProvider): - def __init__(self) -> None: - self.reports: list[MassBalanceData] = [] - - def log(self, report_data: MassBalanceData) -> None: - self.reports.append(report_data) - - def test_float32_rain_rate_representation_error_is_negligible(): """Storing 10 mm/h as float32 should not cause a meaningful daily error.""" stored_rate = np.float32(RAIN_RATE) @@ -122,7 +114,7 @@ def _run_rain_only_simulation( dtinf=timestep, infiltration_model=InfiltrationModelType.NULL, ) - mass_balance_output = CaptureMassBalanceOutputProvider() + mass_balance_output = MemoryMassBalanceOutputProvider() simulation = ( SimulationBuilder(config, np.zeros(shape, dtype=np.bool_), dtype) .with_domain_data(domain_data) @@ -174,7 +166,7 @@ def _run_combined_removal_simulation( dtinf=1000.0, infiltration_model=infiltration_model, ) - mass_balance_output = CaptureMassBalanceOutputProvider() + mass_balance_output = MemoryMassBalanceOutputProvider() raster_output = MemoryRasterOutputProvider() dtype = np.float32 simulation = (