feat: update assessment portal

tests/test_prediction.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
 
 import pandas as pd
 
@@ -13,9 +15,16 @@ from app.prediction import (
     interpolate_probability,
     secure_upload_name,
     validate_input_frame,
+    write_prediction_workbook,
 )
 
 
+class DummyCurve:
+    def __init__(self, x: list[float], y: list[float]) -> None:
+        self.x = x
+        self.y = y
+
+
 class PredictionHelpersTest(unittest.TestCase):
     def test_secure_upload_name_accepts_chinese_filename(self) -> None:
         filename, suffix = secure_upload_name("管道数据.xlsx", "run123")
@@ -44,6 +53,29 @@ class PredictionHelpersTest(unittest.TestCase):
 
         self.assertIn("缺少必要字段", ctx.exception.message)
 
+    def test_prediction_workbook_keeps_sample_data_in_one_sheet(self) -> None:
+        curves = [
+            DummyCurve([1, 2], [0.9, 0.7]),
+            DummyCurve([1, 2], [0.8, 0.6]),
+        ]
+        summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}]
+        summary_sheet_rows = [
+            {"管道编号": "P001", "健康概率": 0.7},
+            {"管道编号": "P002", "健康概率": 0.6},
+        ]
+
+        with TemporaryDirectory() as temp_dir:
+            output_path = Path(temp_dir) / "prediction.xlsx"
+            write_prediction_workbook(output_path, curves, summary_rows, summary_sheet_rows)
+
+            workbook = pd.ExcelFile(output_path)
+            self.assertEqual(workbook.sheet_names, ["结果摘要", "样本数据"])
+
+            sample_data = pd.read_excel(output_path, sheet_name="样本数据")
+            self.assertEqual(len(sample_data), 4)
+            self.assertEqual(sample_data["管道编号"].tolist(), ["P001", "P001", "P002", "P002"])
+            self.assertIn("风险概率", sample_data.columns)
+
 
 if __name__ == "__main__":
     unittest.main()