add MemoryMassBalanceOutputProvider

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."""
 

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

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 = (