From 6f9c94a4dd659216feb56103d2a850005d5b764b Mon Sep 17 00:00:00 2001
From: Huarch <huarch97@outlook.com>
Date: Thu, 16 Jul 2026 14:50:22 +0800
Subject: [PATCH] fix(api): include simulation burst ids

---
 app/services/burst_location.py            | 31 +++++++++++++++++++++++
 tests/unit/test_burst_location_service.py | 17 +++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/app/services/burst_location.py b/app/services/burst_location.py
index a1bcf73..049637b 100644
--- a/app/services/burst_location.py
+++ b/app/services/burst_location.py
@@ -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],
diff --git a/tests/unit/test_burst_location_service.py b/tests/unit/test_burst_location_service.py
index 4ca9af5..1734db0 100644
--- a/tests/unit/test_burst_location_service.py
+++ b/tests/unit/test_burst_location_service.py
@@ -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}