fix(api): include simulation burst ids

This commit is contained in:
2026-07-16 14:50:22 +08:00
parent 775ecb8a58
commit ca1579dcc2
2 changed files with 48 additions and 0 deletions
+31
View File
@@ -11,6 +11,7 @@ from app.infra.db.timescaledb.internal_queries import InternalQueries
from app.services.scheme_management import (
query_burst_location_scheme_detail,
query_burst_location_schemes,
query_scheme_list,
scheme_name_exists,
store_scheme_info,
)
@@ -353,9 +354,15 @@ def run_burst_location_by_network(
}
)
if normalized_data_source == "simulation":
simulation_burst_ids = _get_simulation_scheme_burst_ids(
network=network,
scheme_name=simulation_scheme_name,
scheme_type=resolved_simulation_scheme_type,
)
payload["simulation_scheme"] = {
"name": simulation_scheme_name,
"type": resolved_simulation_scheme_type,
"burst_ids": simulation_burst_ids,
}
if scheme_name:
_store_burst_scheme(
@@ -464,6 +471,30 @@ def _validate_time_window(
return start_dt, end_dt
def _get_simulation_scheme_burst_ids(
*, network: str, scheme_name: str | None, scheme_type: str
) -> list[str]:
if not scheme_name:
return []
rows = query_scheme_list(network) or []
for row in rows:
if len(row) < 7:
continue
if row[1] != scheme_name or row[2] != scheme_type:
continue
detail = row[6] if isinstance(row[6], dict) else {}
return _normalize_burst_ids(detail.get("burst_ID"))
return []
def _normalize_burst_ids(value: Any) -> list[str]:
if value is None:
return []
if isinstance(value, (list, tuple, set)):
return _dedupe_ids([str(item) for item in value])
return _dedupe_ids([str(value)])
def _align_observed_series_pair(
*,
ids: list[str],
+17
View File
@@ -75,6 +75,7 @@ def _load_burst_location_module():
scheme_management_module = types.ModuleType("app.services.scheme_management")
scheme_management_module.query_burst_location_scheme_detail = lambda *args, **kwargs: {}
scheme_management_module.query_burst_location_schemes = lambda *args, **kwargs: []
scheme_management_module.query_scheme_list = lambda *args, **kwargs: []
scheme_management_module.scheme_name_exists = lambda *args, **kwargs: False
scheme_management_module.store_scheme_info = lambda *args, **kwargs: None
sys.modules["app.services.scheme_management"] = scheme_management_module
@@ -177,6 +178,21 @@ def test_run_burst_location_uses_single_timerange_with_burst_source_split(monkey
"query_realtime_simulation_by_ids_timerange",
staticmethod(fake_realtime_query),
)
monkeypatch.setattr(
module,
"query_scheme_list",
lambda name: [
(
1,
"BurstSchemeA",
"burst_analysis",
"testuser",
None,
None,
{"burst_ID": ["Pipe-009", "Pipe-010"]},
)
],
)
result = module.run_burst_location_by_network(
network="tjwater",
@@ -194,6 +210,7 @@ def test_run_burst_location_uses_single_timerange_with_burst_source_split(monkey
assert result["simulation_scheme"] == {
"name": "BurstSchemeA",
"type": "burst_analysis",
"burst_ids": ["Pipe-009", "Pipe-010"],
}
assert result["pressure_samples"] == {"burst": 4, "normal": 4}
assert result["flow_samples"] == {"burst": 4, "normal": 4}