TJWaterServerBinary/tests/unit/test_scheme_simulation_timestep.py

import inspect
import json
from datetime import timedelta

import pytest

from app.infra.db.timescaledb.repositories.scheme import SchemeRepository
from app.services.time_api import parse_utc_time


def test_run_simulation_exposes_explicit_valve_control():
    from app.services import simulation

    parameters = inspect.signature(simulation.run_simulation).parameters

    assert "valve_control" in parameters


def test_apply_valve_control_matches_run_simulation_ex_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"] == {
        "link": "V-status",
        "status": "ACTIVE",
        "setting": 1.0,
    }
    assert updates["V-setting"] == {
        "link": "V-setting",
        "status": "OPEN",
        "setting": 2.5,
    }
    assert updates["V-closed"] == {
        "link": "V-closed",
        "status": "CLOSED",
        "setting": 9.0,
    }
    assert updates["V-k"] == {
        "link": "V-k",
        "status": "ACTIVE",
        "setting": 0.1036 * pow(0.5, -3.105),
    }


def _node_result(periods: int) -> list[dict]:
    return [
        {
            "node": "J1",
            "result": [
                {"demand": index, "head": index, "pressure": index, "quality": index}
                for index in range(periods)
            ],
        }
    ]


def _link_result(periods: int) -> list[dict]:
    return [
        {
            "link": "P1",
            "result": [
                {
                    "flow": index,
                    "friction": index,
                    "headloss": index,
                    "quality": index,
                    "reaction": index,
                    "setting": index,
                    "status": index,
                    "velocity": index,
                }
                for index in range(periods)
            ],
        }
    ]


def test_store_scheme_simulation_uses_15_minute_report_step(monkeypatch):
    inserted: dict[str, list[dict]] = {}

    monkeypatch.setattr(
        SchemeRepository,
        "insert_nodes_batch_sync",
        staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
    )
    monkeypatch.setattr(
        SchemeRepository,
        "insert_links_batch_sync",
        staticmethod(lambda conn, data: inserted.setdefault("links", data)),
    )

    SchemeRepository.store_scheme_simulation_result_sync(
        conn=object(),
        scheme_type="burst_analysis",
        scheme_name="five_hour_case",
        node_result_list=_node_result(21),
        link_result_list=_link_result(21),
        result_start_time="2026-07-16T00:00:00Z",
        num_periods=21,
        result_timestep_seconds=900,
    )

    start_time = parse_utc_time("2026-07-16T00:00:00Z")
    assert len(inserted["nodes"]) == 21
    assert inserted["nodes"][0]["time"] == start_time
    assert inserted["nodes"][-1]["time"] == start_time + timedelta(hours=5)
    assert inserted["links"][-1]["time"] == start_time + timedelta(hours=5)


def test_store_scheme_simulation_uses_hourly_report_step(monkeypatch):
    inserted: dict[str, list[dict]] = {}

    monkeypatch.setattr(
        SchemeRepository,
        "insert_nodes_batch_sync",
        staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
    )
    monkeypatch.setattr(
        SchemeRepository,
        "insert_links_batch_sync",
        staticmethod(lambda conn, data: inserted.setdefault("links", data)),
    )

    SchemeRepository.store_scheme_simulation_result_sync(
        conn=object(),
        scheme_type="burst_analysis",
        scheme_name="hourly_case",
        node_result_list=_node_result(6),
        link_result_list=_link_result(6),
        result_start_time="2026-07-16T00:00:00Z",
        num_periods=6,
        result_timestep_seconds=3600,
    )

    start_time = parse_utc_time("2026-07-16T00:00:00Z")
    assert [item["time"] for item in inserted["nodes"]] == [
        start_time + timedelta(hours=index) for index in range(6)
    ]


def test_run_simulation_passes_report_step_for_extended_scheme(monkeypatch):
    import app.services.simulation as simulation

    time_updates: list[dict] = []
    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: time_updates.append(changeset.operations[0]),
    )
    monkeypatch.setattr(simulation, "run_project", lambda name: json.dumps({
        "simulation_result": "successful",
        "output": {
            "times": {"num_periods": 21, "report_step": 900},
            "node_results": _node_result(21),
            "link_results": _link_result(21),
        },
    }))
    monkeypatch.setattr(
        simulation.TimescaleInternalStorage,
        "store_scheme_simulation",
        staticmethod(lambda *args, **kwargs: storage_calls.append((args, kwargs))),
    )

    simulation.run_simulation(
        name="fengyang",
        simulation_type="extended",
        modify_pattern_start_time="2026-07-16T00:00:00+08:00",
        modify_total_duration=18000,
        scheme_type="burst_analysis",
        scheme_name="five_hour_case",
    )

    assert time_updates[0]["DURATION"] == "05:00:00"
    assert time_updates[0]["REPORT TIMESTEP"] == "1:00"
    assert storage_calls[0][0][5] == 21
    assert storage_calls[0][0][6] == 900