diff --git a/app/infra/db/timescaledb/internal_queries.py b/app/infra/db/timescaledb/internal_queries.py index c4fcfe9..16d45fe 100644 --- a/app/infra/db/timescaledb/internal_queries.py +++ b/app/infra/db/timescaledb/internal_queries.py @@ -50,6 +50,7 @@ class InternalStorage: link_result_list: List[dict], result_start_time: str, num_periods: int = 1, + result_timestep_seconds: int | None = None, db_name: str = None, max_retries: int = 3, ): @@ -70,6 +71,7 @@ class InternalStorage: link_result_list, result_start_time, num_periods, + result_timestep_seconds, ) break # 成功 except Exception as e: diff --git a/app/infra/db/timescaledb/repositories/scheme.py b/app/infra/db/timescaledb/repositories/scheme.py index f0960b3..3903fc0 100644 --- a/app/infra/db/timescaledb/repositories/scheme.py +++ b/app/infra/db/timescaledb/repositories/scheme.py @@ -3,10 +3,24 @@ from datetime import datetime, timedelta from collections import defaultdict from psycopg import AsyncConnection, Connection, sql import app.services.globals as globals -from app.services.time_api import parse_utc_time +from app.services.time_api import parse_clock_duration_seconds, parse_utc_time class SchemeRepository: + @staticmethod + def _get_result_timestep(result_timestep_seconds: int | None) -> timedelta: + if result_timestep_seconds is not None: + if result_timestep_seconds <= 0: + raise ValueError("result_timestep_seconds must be greater than 0.") + return timedelta(seconds=result_timestep_seconds) + + timestep_seconds = parse_clock_duration_seconds( + globals.hydraulic_timestep, + field_name="HYDRAULIC TIMESTEP", + ) + if timestep_seconds <= 0: + raise ValueError("HYDRAULIC TIMESTEP must be greater than 0.") + return timedelta(seconds=timestep_seconds) # --- Link Simulation --- @@ -452,6 +466,7 @@ class SchemeRepository: link_result_list: List[Dict[str, any]], result_start_time: str, num_periods: int = 1, + result_timestep_seconds: int | None = None, ): """ Store scheme simulation results to TimescaleDB. @@ -468,20 +483,16 @@ class SchemeRepository: result_start_time, field_name="result_start_time" ) - timestep_parts = globals.hydraulic_timestep.split(":") - timestep = timedelta( - hours=int(timestep_parts[0]), - minutes=int(timestep_parts[1]), - seconds=int(timestep_parts[2]), - ) + timestep = SchemeRepository._get_result_timestep(result_timestep_seconds) # Prepare node data for batch insert node_data = [] for node_result in node_result_list: node_id = node_result.get("node") - for period_index in range(num_periods): + result_rows = node_result.get("result", []) + for period_index in range(min(num_periods, len(result_rows))): current_time = simulation_time + (timestep * period_index) - data = node_result.get("result", [])[period_index] + data = result_rows[period_index] node_data.append( { "time": current_time, @@ -499,9 +510,10 @@ class SchemeRepository: link_data = [] for link_result in link_result_list: link_id = link_result.get("link") - for period_index in range(num_periods): + result_rows = link_result.get("result", []) + for period_index in range(min(num_periods, len(result_rows))): current_time = simulation_time + (timestep * period_index) - data = link_result.get("result", [])[period_index] + data = result_rows[period_index] link_data.append( { "time": current_time, @@ -535,6 +547,7 @@ class SchemeRepository: link_result_list: List[Dict[str, any]], result_start_time: str, num_periods: int = 1, + result_timestep_seconds: int | None = None, ): """ Store scheme simulation results to TimescaleDB (sync version). @@ -551,20 +564,16 @@ class SchemeRepository: result_start_time, field_name="result_start_time" ) - timestep_parts = globals.hydraulic_timestep.split(":") - timestep = timedelta( - hours=int(timestep_parts[0]), - minutes=int(timestep_parts[1]), - seconds=int(timestep_parts[2]), - ) + timestep = SchemeRepository._get_result_timestep(result_timestep_seconds) # Prepare node data for batch insert node_data = [] for node_result in node_result_list: node_id = node_result.get("node") - for period_index in range(num_periods): + result_rows = node_result.get("result", []) + for period_index in range(min(num_periods, len(result_rows))): current_time = simulation_time + (timestep * period_index) - data = node_result.get("result", [])[period_index] + data = result_rows[period_index] node_data.append( { "time": current_time, @@ -582,9 +591,10 @@ class SchemeRepository: link_data = [] for link_result in link_result_list: link_id = link_result.get("link") - for period_index in range(num_periods): + result_rows = link_result.get("result", []) + for period_index in range(min(num_periods, len(result_rows))): current_time = simulation_time + (timestep * period_index) - data = link_result.get("result", [])[period_index] + data = result_rows[period_index] link_data.append( { "time": current_time, diff --git a/app/services/simulation.py b/app/services/simulation.py index 5e910e9..f567930 100644 --- a/app/services/simulation.py +++ b/app/services/simulation.py @@ -1260,6 +1260,9 @@ def run_simulation( node_result, link_result, modify_pattern_start_time, db_name=db_name ) elif simulation_type.upper() == "EXTENDED": + result_timestep_seconds = times_info.get("report_step") + if result_timestep_seconds is None: + raise RuntimeError("run_project output missing times.report_step") TimescaleInternalStorage.store_scheme_simulation( scheme_type, scheme_name, @@ -1267,6 +1270,7 @@ def run_simulation( link_result, modify_pattern_start_time, num_periods_result, + result_timestep_seconds, db_name=db_name, ) endtime = time.time() diff --git a/tests/unit/test_scheme_simulation_timestep.py b/tests/unit/test_scheme_simulation_timestep.py new file mode 100644 index 0000000..72c37c6 --- /dev/null +++ b/tests/unit/test_scheme_simulation_timestep.py @@ -0,0 +1,154 @@ +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 _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