TJWaterServerBinary/tests/api/test_simulation_endpoints.py

from datetime import datetime, timezone
from types import SimpleNamespace

from fastapi.testclient import TestClient

from tests.conftest import build_test_app, install_stub, load_module_from_path


def _load_simulation_module(monkeypatch):
    install_stub(monkeypatch, "app.services", package=True)
    def parse_aware_time(value, field_name="datetime"):
        dt = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
        if dt.tzinfo is None:
            raise ValueError(f"{field_name} is missing timezone information.")
        return dt

    def parse_utc_time(value, field_name="datetime"):
        return parse_aware_time(value, field_name=field_name).astimezone(
            timezone.utc
        )

    def parse_clock_duration_seconds(value, field_name="duration"):
        parts = [int(part) for part in value.split(":")]
        hours, minutes = parts[0], parts[1]
        seconds = parts[2] if len(parts) == 3 else 0
        return hours * 3600 + minutes * 60 + seconds

    install_stub(
        monkeypatch,
        "app.services.simulation",
        {
            "get_time": lambda name: {"HYDRAULIC TIMESTEP": "0:15:00"},
            "run_simulation": lambda **kwargs: None,
            "query_corresponding_element_id_and_query_id": lambda name: SimpleNamespace(
                fixed_pumps={}, variable_pumps={}
            ),
            "query_corresponding_pattern_id_and_query_id": lambda name: None,
            "query_non_realtime_region": lambda name: [],
            "get_source_outflow_region_id": lambda name, region_result: {},
            "query_realtime_region_pipe_flow_and_demand_id": lambda name, region_result: {},
            "query_pipe_flow_region_patterns": lambda name: {},
            "query_non_realtime_region_patterns": lambda name, region_result: {},
            "get_realtime_region_patterns": lambda name, source_outflow_region_id, realtime_region_pipe_flow_and_demand_id: ({}, {}),
        },
    )
    install_stub(
        monkeypatch,
        "app.services.tjnetwork",
        {
            "run_project": lambda network: "report",
            "run_project_return_dict": lambda network: {"output": {}, "report": "ok"},
            "run_inp": lambda network: "inp-report",
            "dump_output": lambda output: f"dump::{output}",
        },
    )
    install_stub(
        monkeypatch,
        "app.services.simulation_scenarios",
        {
            "burst_analysis": lambda *args, **kwargs: "burst",
            "valve_close_analysis": lambda *args, **kwargs: "valve",
            "flushing_analysis": lambda *args, **kwargs: "flush",
            "contaminant_simulation": lambda *args, **kwargs: "contaminant",
            "pressure_regulation": lambda *args, **kwargs: "pressure",
        },
    )
    install_stub(
        monkeypatch,
        "app.services.network_import",
        {"network_update": lambda *args, **kwargs: "updated"},
    )
    install_stub(
        monkeypatch,
        "app.services.valve_isolation",
        {"analyze_valve_isolation": lambda *args, **kwargs: {}},
    )
    return load_module_from_path(
        "tests_simulation_endpoints_module",
        "app/api/v1/endpoints/simulation.py",
    )


def _build_authenticated_client(module) -> TestClient:
    app = build_test_app(module.router, "/api/v1")
    app.dependency_overrides[module.get_current_keycloak_username] = lambda: "alice"
    return TestClient(app)


def test_run_project_endpoint_returns_plain_text(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    monkeypatch.setattr(module, "run_project", lambda network: f"report::{network}")
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post("/api/v1/project-runs", params={"network": "demo"})

    assert response.status_code == 200
    assert response.text == "report::demo"


def test_removed_legacy_simulation_workflows_are_not_exposed(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    for path in (
        "/api/v1/project-managements",
        "/api/v1/scheduling-analyses",
        "/api/v1/daily-scheduling-analyses",
        "/api/v1/water-age-analyses",
        "/api/v1/inp-runs",
        "/api/v1/outputs",
        "/api/v1/pump-failure-events",
    ):
        assert client.post(path).status_code == 404


def test_removed_basic_pressure_regulation_is_not_exposed(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post(
        "/api/v1/pressure-regulation-calculations",
        params={
            "network": "demo",
            "target_node": "J1",
            "target_pressure": 50,
        },
    )

    assert response.status_code == 404


def test_pressure_regulation_requires_scheme_name(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post(
        "/api/v1/pressure-regulation-analyses",
        json={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "pump_control": {},
        },
    )

    assert response.status_code == 422


def test_run_simulation_manually_by_date_uses_utc_aware_timestamps(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured_calls = []

    monkeypatch.setattr(
        module.simulation,
        "run_simulation",
        lambda **kwargs: captured_calls.append(kwargs),
    )

    module.run_simulation_manually_by_date(
        "demo",
        datetime(2025, 1, 1, 19, 4, 5, tzinfo=timezone.utc),
        30,
    )

    assert [call["modify_pattern_start_time"] for call in captured_calls] == [
        "2025-01-01T19:04:05+00:00",
        "2025-01-01T19:19:05+00:00",
    ]


def test_run_simulation_manually_by_date_uses_hydraulic_timestep(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured_calls = []

    monkeypatch.setattr(
        module.simulation,
        "get_time",
        lambda name: {"HYDRAULIC TIMESTEP": "1:00"},
    )
    monkeypatch.setattr(
        module.simulation,
        "run_simulation",
        lambda **kwargs: captured_calls.append(kwargs),
    )

    module.run_simulation_manually_by_date(
        "demo",
        datetime(2025, 1, 1, 16, 0, 0, tzinfo=timezone.utc),
        120,
    )

    assert [call["modify_pattern_start_time"] for call in captured_calls] == [
        "2025-01-01T16:00:00+00:00",
        "2025-01-01T17:00:00+00:00",
    ]


def test_runsimulationmanuallybydate_endpoint_accepts_timezone_aware_start_time(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_run(network_name, start_time, duration):
        captured["network_name"] = network_name
        captured["start_time"] = start_time
        captured["duration"] = duration

    monkeypatch.setattr(module, "run_simulation_manually_by_date", fake_run)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post(
        "/api/v1/simulation-runs",
        json={
            "name": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "duration": 30,
        },
    )

    assert response.status_code == 200
    assert response.json() == {"status": "success"}
    assert captured["network_name"] == "demo"
    assert captured["duration"] == 30
    assert captured["start_time"].isoformat() == "2025-01-01T19:04:05+00:00"


def test_runsimulationmanuallybydate_endpoint_rejects_naive_start_time(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post(
        "/api/v1/simulation-runs",
        json={
            "name": "demo",
            "start_time": "2025-01-02T03:04:05",
            "duration": 30,
        },
    )

    assert response.status_code == 422


def test_valve_close_endpoint_passes_scheme_name(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_valve_close_analysis(**kwargs):
        captured.update(kwargs)
        return "ok"

    monkeypatch.setattr(module, "valve_close_analysis", fake_valve_close_analysis)
    client = TestClient(build_test_app(module.router, "/api/v1"))

    response = client.post(
        "/api/v1/valve-closure-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "valves": ["V1", "V2"],
            "duration": 900,
            "scheme_name": "valve_case_01",
        },
    )

    assert response.status_code == 200
    assert response.text == "ok"
    assert captured == {
        "name": "demo",
        "modify_pattern_start_time": "2025-01-02T03:04:05+08:00",
        "modify_total_duration": 900,
        "modify_valve_opening": {"V1": 0.0, "V2": 0.0},
        "scheme_name": "valve_case_01",
    }


def test_burst_endpoint_passes_current_username(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_burst_analysis(**kwargs):
        captured.update(kwargs)

    monkeypatch.setattr(module, "burst_analysis", fake_burst_analysis)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/burst-analyses",
        params={
            "network": "demo",
            "modify_pattern_start_time": "2025-01-02T03:04:05+08:00",
            "burst_ID": ["P1"],
            "burst_size": [10.0],
            "modify_total_duration": 900,
            "scheme_name": "burst_case_01",
        },
    )

    assert response.status_code == 200
    assert response.text == '"success"'
    assert captured["username"] == "alice"


def test_burst_endpoint_rejects_mismatched_pipe_and_size_counts(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    called = False

    def fake_burst_analysis(**kwargs):
        nonlocal called
        called = True

    monkeypatch.setattr(module, "burst_analysis", fake_burst_analysis)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/burst-analyses",
        params=[
            ("network", "demo"),
            ("modify_pattern_start_time", "2025-01-02T03:04:05+08:00"),
            ("burst_ID", "P1"),
            ("burst_ID", "P2"),
            ("burst_size", "10.0"),
            ("modify_total_duration", "900"),
            ("scheme_name", "burst_case_01"),
        ],
    )

    assert response.status_code == 422
    assert called is False


def test_flushing_endpoint_passes_required_scheme_name(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_flushing_analysis(**kwargs):
        captured.update(kwargs)
        return "ok"

    monkeypatch.setattr(module, "flushing_analysis", fake_flushing_analysis)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "valves": ["V1"],
            "valves_k": [0.5],
            "drainage_node_ID": "N1",
            "flush_flow": 100.0,
            "duration": 900,
            "scheme_name": "flush_case_01",
        },
    )

    assert response.status_code == 200
    assert response.text == "ok"
    assert captured == {
        "name": "demo",
        "modify_pattern_start_time": "2025-01-02T03:04:05+08:00",
        "modify_total_duration": 900,
        "modify_valve_opening": {"V1": 0.5},
        "valve_control": None,
        "drainage_node_ID": "N1",
        "flushing_flow": 100.0,
        "scheme_name": "flush_case_01",
        "username": "alice",
    }


def test_flushing_endpoint_allows_omitting_valves(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_flushing_analysis(**kwargs):
        captured.update(kwargs)
        return "ok"

    monkeypatch.setattr(module, "flushing_analysis", fake_flushing_analysis)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "drainage_node_ID": "N1",
            "scheme_name": "flush_without_valves",
        },
    )

    assert response.status_code == 200
    assert response.text == "ok"
    assert captured["modify_valve_opening"] is None
    assert captured["valve_control"] is None
    assert captured["drainage_node_ID"] == "N1"


def test_flushing_endpoint_passes_explicit_valve_control(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_flushing_analysis(**kwargs):
        captured.update(kwargs)
        return "ok"

    monkeypatch.setattr(module, "flushing_analysis", fake_flushing_analysis)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params=[
            ("network", "demo"),
            ("start_time", "2025-01-02T03:04:05+08:00"),
            ("valves", "V1"),
            ("valves", "V2"),
            ("valve_statuses", "ACTIVE"),
            ("valve_statuses", "CLOSED"),
            ("valve_settings", "2.5"),
            ("valve_settings", ""),
            ("drainage_node_ID", "N1"),
            ("scheme_name", "flush_with_valve_control"),
        ],
    )

    assert response.status_code == 200
    assert captured["modify_valve_opening"] is None
    assert captured["valve_control"] == {
        "V1": {"status": "ACTIVE", "setting": "2.5"},
        "V2": {"status": "CLOSED"},
    }


def test_flushing_endpoint_requires_setting_for_active_valve(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "valves": "V1",
            "valve_statuses": "ACTIVE",
            "drainage_node_ID": "N1",
            "scheme_name": "flush_without_active_setting",
        },
    )

    assert response.status_code == 422


def test_flushing_endpoint_rejects_mixed_valve_control_modes(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "valves": "V1",
            "valves_k": 0.5,
            "valve_statuses": "ACTIVE",
            "valve_settings": "2.5",
            "drainage_node_ID": "N1",
            "scheme_name": "flush_with_mixed_controls",
        },
    )

    assert response.status_code == 422


def test_flushing_endpoint_rejects_settings_without_statuses(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "valves": "V1",
            "valves_k": 0.5,
            "valve_settings": "2.5",
            "drainage_node_ID": "N1",
            "scheme_name": "flush_with_orphan_settings",
        },
    )

    assert response.status_code == 422


def test_flushing_endpoint_requires_drainage_node(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/flushing-analyses",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "scheme_name": "flush_without_drainage_node",
        },
    )

    assert response.status_code == 422


def test_contaminant_endpoint_passes_current_username(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    captured = {}

    def fake_contaminant_simulation(**kwargs):
        captured.update(kwargs)
        return "ok"

    monkeypatch.setattr(module, "contaminant_simulation", fake_contaminant_simulation)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/contaminant-simulations",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "source": "N1",
            "concentration": 10.0,
            "duration": 900,
            "scheme_name": "contaminant_case_01",
        },
    )

    assert response.status_code == 200
    assert response.text == "ok"
    assert captured["username"] == "alice"


def test_contaminant_endpoint_requires_scheme_name(monkeypatch):
    module = _load_simulation_module(monkeypatch)
    client = _build_authenticated_client(module)

    response = client.post(
        "/api/v1/contaminant-simulations",
        params={
            "network": "demo",
            "start_time": "2025-01-02T03:04:05+08:00",
            "source": "N1",
            "concentration": 10.0,
            "duration": 900,
        },
    )

    assert response.status_code == 422