TJWaterServerBinary/tests/unit/test_analysis_simulation.py

import inspect
import json
from uuid import uuid4


def test_run_simulation_exposes_explicit_valve_control():
    from app.services import simulation

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


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_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)
    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",
    )

    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"


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

    run_id = uuid4()
    lifecycle_calls: list[tuple] = []
    monkeypatch.setattr(simulation, "open_project", 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",
        )

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