refactor(db)!: finalize pooled WNDB v2 migration

tests/unit/test_analysis_simulation.py

@@ -1,7 +1,11 @@
 import inspect
 import json
+from contextlib import contextmanager
+from unittest.mock import Mock
 from uuid import uuid4
 
+import pytest
+
 
 def test_run_simulation_exposes_explicit_valve_control():
     from app.services import simulation
@@ -9,6 +13,34 @@ def test_run_simulation_exposes_explicit_valve_control():
     assert "valve_control" in inspect.signature(simulation.run_simulation).parameters
 
 
+def test_extended_runner_cleans_temporary_database_after_failure(monkeypatch):
+    from app.algorithms.simulation import runner
+
+    lifecycle: list[tuple[str, str]] = []
+
+    @contextmanager
+    def temporary_project(project: str, purpose: str):
+        lifecycle.append(("create", project))
+        try:
+            yield "isolated_project"
+        finally:
+            lifecycle.append(("delete", project))
+
+    monkeypatch.setattr(runner, "temporary_project_database", temporary_project)
+
+    @runner._clean_extended_simulation
+    def fail(name, simulation_type, *, _temporary_project=None):
+        assert name == "demo"
+        assert simulation_type == "extended"
+        assert _temporary_project == "isolated_project"
+        raise RuntimeError("simulation failed")
+
+    with pytest.raises(RuntimeError, match="simulation failed"):
+        fail("demo", "extended")
+
+    assert lifecycle == [("create", "demo"), ("delete", "demo")]
+
+
 def test_apply_valve_control_matches_runner_semantics(monkeypatch):
     from app.services import simulation
 
@@ -46,12 +78,58 @@ def test_apply_valve_control_matches_runner_semantics(monkeypatch):
     assert updates["V-k"]["setting"] == 0.1036 * pow(0.5, -3.105)
 
 
+def test_primary_demand_update_preserves_additional_categories():
+    from app.services import simulation
+
+    demand_set = {
+        "junction": "J1",
+        "demands": [
+            {"demand": 1.0, "pattern": "P1", "category": "domestic"},
+            {"demand": 2.0, "pattern": "P2", "category": "industrial"},
+        ],
+    }
+
+    simulation._primary_demand(demand_set)["demand"] = 3.0
+
+    assert demand_set["demands"] == [
+        {"demand": 3.0, "pattern": "P1", "category": "domestic"},
+        {"demand": 2.0, "pattern": "P2", "category": "industrial"},
+    ]
+    assert simulation._primary_demand_pattern(demand_set) == "P1"
+
+
+def test_primary_demand_is_created_for_empty_junction():
+    from app.services import simulation
+
+    demand_set = {"junction": "J1", "demands": []}
+
+    primary = simulation._primary_demand(demand_set)
+
+    assert primary == {"demand": 0.0, "pattern": None, "category": None}
+    with pytest.raises(ValueError, match="has no demand pattern"):
+        simulation._primary_demand_pattern(demand_set)
+
+
 def test_extended_simulation_stores_results_by_run_id(monkeypatch):
     from app.services import simulation
 
     run_id = uuid4()
     storage_calls: list[tuple] = []
-    monkeypatch.setattr(simulation, "open_project", lambda name: None)
+    transaction_calls: list[tuple[str, str]] = []
+
+    @contextmanager
+    def project_transaction(name):
+        transaction_calls.append(("begin", name))
+        try:
+            yield object()
+        finally:
+            transaction_calls.append(("end", name))
+
+    refresh_mock = Mock()
+    monkeypatch.setattr(simulation, "project_transaction", project_transaction)
+    monkeypatch.setattr(
+        simulation, "refresh_materialized_views_after_commit", refresh_mock
+    )
     monkeypatch.setattr(
         simulation,
         "get_time",
@@ -104,6 +182,8 @@ def test_extended_simulation_stores_results_by_run_id(monkeypatch):
     assert kwargs["db_name"] == "demo"
     assert returned_run_id == run_id
     assert lifecycle_calls[-1][1]["status"] == "completed"
+    assert transaction_calls == [("begin", "demo"), ("end", "demo")]
+    refresh_mock.assert_called_once_with("demo")
 
 
 def test_extended_simulation_marks_run_failed_when_result_storage_fails(monkeypatch):
@@ -111,7 +191,15 @@ def test_extended_simulation_marks_run_failed_when_result_storage_fails(monkeypa
 
     run_id = uuid4()
     lifecycle_calls: list[tuple] = []
-    monkeypatch.setattr(simulation, "open_project", lambda name: None)
+
+    @contextmanager
+    def project_transaction(_name):
+        yield object()
+
+    monkeypatch.setattr(simulation, "project_transaction", project_transaction)
+    monkeypatch.setattr(
+        simulation, "refresh_materialized_views_after_commit", lambda _name: None
+    )
     monkeypatch.setattr(
         simulation,
         "get_time",