TJWaterServerBinary/tests/unit/test_analysis_simulation.py

import inspect
import json
from contextlib import contextmanager
from unittest.mock import Mock
from uuid import uuid4

import pytest


def _empty_scada_mappings(simulation):
    return simulation.ScadaElementMappings(
        reservoirs={},
        tanks={},
        fixed_pumps={},
        variable_pumps={},
        pressure={},
        demand={},
        quality={},
    )


def test_run_simulation_accepts_explicit_valve_control():
    from app.services import simulation

    assert "valve_control" in inspect.signature(simulation.run_simulation).parameters


def test_valve_close_analysis_uses_normalized_scheme_type(monkeypatch):
    from app.services import simulation_scenarios

    captured = {}
    monkeypatch.setattr(simulation_scenarios, "get_option", lambda _name: {})
    monkeypatch.setattr(
        simulation_scenarios,
        "set_option",
        lambda _name, _changes: None,
    )
    monkeypatch.setattr(
        simulation_scenarios.simulation,
        "run_simulation",
        lambda **kwargs: captured.update(kwargs),
    )

    simulation_scenarios.valve_close_analysis.__wrapped__(
        name="demo",
        modify_pattern_start_time="2026-01-01T00:00:00+08:00",
        modify_valve_opening={"V1": 0.0},
        scheme_name="valve_case",
        _temporary_project="temporary_demo",
    )

    assert captured["scheme_type"] == "valve_close_analysis"
    assert captured["result_db_name"] == "demo"


def test_apply_valve_control_matches_runner_semantics(monkeypatch):
    from app.services import simulation

    updates: dict[str, dict] = {}
    monkeypatch.setattr(
        simulation,
        "get_status",
        lambda project_name, valve_name: {
            "link": valve_name,
            "status": "OPEN",
            "setting": 1.0,
        },
    )
    monkeypatch.setattr(
        simulation,
        "set_status",
        lambda project_name, changeset: updates.update(
            {changeset.operations[0]["link"]: changeset.operations[0].copy()}
        ),
    )

    simulation._apply_valve_control(
        "demo",
        {
            "V-status": {"status": "ACTIVE"},
            "V-setting": {"setting": 2.5},
            "V-closed": {"status": "ACTIVE", "setting": 9.0, "k": 0},
            "V-k": {"status": "ACTIVE", "setting": 9.0, "k": 0.5},
        },
    )

    assert updates["V-status"]["status"] == "ACTIVE"
    assert updates["V-setting"]["setting"] == 2.5
    assert updates["V-closed"]["status"] == "CLOSED"
    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] = []
    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",
        lambda name: {
            "HYDRAULIC TIMESTEP": "00:15:00",
            "REPORT TIMESTEP": "1:00",
            "DURATION": "0:00",
            "PATTERN START": "0:00",
        },
    )
    monkeypatch.setattr(simulation, "set_time", lambda name, changeset: None)
    monkeypatch.setattr(
        simulation,
        "run_project",
        lambda name: json.dumps(
            {
                "output": {
                    "times": {"num_periods": 2, "report_step": 900},
                    "node_results": [{"node": "J1", "result": [{}, {}]}],
                    "link_results": [{"link": "P1", "result": [{}, {}]}],
                }
            }
        ),
    )
    lifecycle_calls: list[tuple] = []
    monkeypatch.setattr(simulation, "create_analysis_run", lambda **kwargs: run_id)
    monkeypatch.setattr(
        simulation,
        "update_analysis_run",
        lambda *args, **kwargs: lifecycle_calls.append((args, kwargs)),
    )
    monkeypatch.setattr(
        simulation.TimescaleInternalStorage,
        "store_analysis_simulation",
        staticmethod(lambda *args, **kwargs: storage_calls.append((args, kwargs))),
    )

    returned_run_id = simulation.run_simulation(
        name="demo",
        simulation_type="extended",
        modify_pattern_start_time="2026-07-16T00:00:00+08:00",
        modify_total_duration=900,
        scheme_type="burst_analysis",
        scheme_name="case",
        scada_mappings=_empty_scada_mappings(simulation),
    )

    args, kwargs = storage_calls[0]
    assert args[0] == run_id
    assert args[4:] == (2, 900)
    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):
    from app.services import simulation

    run_id = uuid4()
    lifecycle_calls: list[tuple] = []

    @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",
        lambda name: {
            "HYDRAULIC TIMESTEP": "00:15:00",
            "REPORT TIMESTEP": "1:00",
            "DURATION": "0:00",
            "PATTERN START": "0:00",
        },
    )
    monkeypatch.setattr(simulation, "set_time", lambda name, changeset: None)
    monkeypatch.setattr(
        simulation,
        "run_project",
        lambda name: json.dumps(
            {
                "output": {
                    "times": {"num_periods": 1, "report_step": 900},
                    "node_results": [{"node": "J1", "result": [{}]}],
                    "link_results": [{"link": "P1", "result": [{}]}],
                }
            }
        ),
    )
    monkeypatch.setattr(simulation, "create_analysis_run", lambda **kwargs: run_id)
    monkeypatch.setattr(
        simulation,
        "update_analysis_run",
        lambda *args, **kwargs: lifecycle_calls.append((args, kwargs)),
    )

    def fail_storage(*args, **kwargs):
        raise RuntimeError("timescale write failed")

    monkeypatch.setattr(
        simulation.TimescaleInternalStorage,
        "store_analysis_simulation",
        staticmethod(fail_storage),
    )

    import pytest

    with pytest.raises(RuntimeError, match="timescale write failed"):
        simulation.run_simulation(
            name="demo",
            simulation_type="extended",
            modify_pattern_start_time="2026-07-16T00:00:00+08:00",
            modify_total_duration=900,
            scheme_type="burst_analysis",
            scheme_name="case",
            scada_mappings=_empty_scada_mappings(simulation),
        )

    assert [call[1]["status"] for call in lifecycle_calls] == ["failed"]