feat: update prediction workbook output

tests/test_prediction.py

@@ -3,16 +3,22 @@ from __future__ import annotations
 import unittest
 from pathlib import Path
 from tempfile import TemporaryDirectory
+from unittest.mock import patch
 
 import pandas as pd
+from openpyxl import load_workbook
 
 from app.prediction import (
     FEATURES,
     ID_COLUMN,
+    PIPE_AGE_COLUMN,
     PredictionError,
     estimate_remaining_life,
     grade_info,
     interpolate_probability,
+    prepare_model_features,
+    read_input_file,
+    render_survival_chart,
     secure_upload_name,
     validate_input_frame,
     write_prediction_workbook,
@@ -53,6 +59,37 @@ class PredictionHelpersTest(unittest.TestCase):
 
         self.assertIn("缺少必要字段", ctx.exception.message)
 
+    def test_read_input_file_preserves_text_pipe_ids(self) -> None:
+        with TemporaryDirectory() as temp_dir:
+            output_path = Path(temp_dir) / "input.xlsx"
+            pd.DataFrame({ID_COLUMN: ["00123"], FEATURES[0]: [1]}).to_excel(output_path, index=False)
+
+            df = read_input_file(output_path, ".xlsx")
+
+            self.assertEqual(df[ID_COLUMN].tolist(), ["00123"])
+
+    def test_prepare_model_features_maps_material_aliases_to_codes(self) -> None:
+        rows = []
+        for value in ["镀锌", "2-钢塑", 13]:
+            row = {feature: 1 for feature in FEATURES}
+            row["管材"] = value
+            rows.append(row)
+        df = pd.DataFrame(rows)
+
+        x_test = prepare_model_features(df)
+
+        self.assertEqual(x_test["管材"].tolist(), [1, 2, 13])
+
+    def test_prepare_model_features_rejects_invalid_material_alias(self) -> None:
+        row = {feature: 1 for feature in FEATURES}
+        row["管材"] = "未知管材"
+        df = pd.DataFrame([row])
+
+        with self.assertRaises(PredictionError) as ctx:
+            prepare_model_features(df)
+
+        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]),
@@ -60,8 +97,8 @@ class PredictionHelpersTest(unittest.TestCase):
         ]
         summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}]
         summary_sheet_rows = [
-            {"管道编号": "P001", "健康概率": 0.7},
-            {"管道编号": "P002", "健康概率": 0.6},
+            {"管道编号": "P001", "管龄": "10 年", "健康概率": 0.7, "预计剩余寿命": 2.0, "健康等级": "IV级"},
+            {"管道编号": "P002", "管龄": "12 年", "健康概率": 0.6, "预计剩余寿命": 1.5, "健康等级": "III级"},
         ]
 
         with TemporaryDirectory() as temp_dir:
@@ -71,10 +108,51 @@ class PredictionHelpersTest(unittest.TestCase):
             workbook = pd.ExcelFile(output_path)
             self.assertEqual(workbook.sheet_names, ["结果摘要", "样本数据"])
 
+            summary_data = pd.read_excel(output_path, sheet_name="结果摘要")
+            self.assertEqual(summary_data.columns.tolist(), [ID_COLUMN, PIPE_AGE_COLUMN, "健康等级"])
+
             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)
+            self.assertEqual(sample_data.columns.tolist(), ["管龄（年）", "P001", "P002"])
+            self.assertEqual(len(sample_data), 2)
+            self.assertEqual(sample_data["管龄（年）"].tolist(), [1, 2])
+            self.assertEqual(sample_data["P001"].tolist(), [0.9, 0.7])
+            self.assertEqual(sample_data["P002"].tolist(), [0.8, 0.6])
+
+            openpyxl_workbook = load_workbook(output_path)
+            sample_worksheet = openpyxl_workbook["样本数据"]
+            self.assertEqual(sample_worksheet.freeze_panes, "A2")
+            self.assertEqual(sample_worksheet.auto_filter.ref, "A1:C3")
+            self.assertEqual(sample_worksheet["B2"].number_format, "0.0%")
+
+    def test_prediction_workbook_writes_pipe_ids_as_excel_text(self) -> None:
+        curves = [DummyCurve([1], [0.9])]
+        summary_rows = [{"pipe_id": "00123"}]
+        summary_sheet_rows = [{ID_COLUMN: "00123", PIPE_AGE_COLUMN: "10 年", "健康等级": "V级"}]
+
+        with TemporaryDirectory() as temp_dir:
+            output_path = Path(temp_dir) / "prediction.xlsx"
+            write_prediction_workbook(output_path, curves, summary_rows, summary_sheet_rows)
+
+            workbook = load_workbook(output_path)
+            summary_worksheet = workbook["结果摘要"]
+            self.assertEqual(summary_worksheet["A2"].value, "00123")
+            self.assertEqual(summary_worksheet["A2"].data_type, "s")
+            self.assertEqual(summary_worksheet["A2"].number_format, "@")
+
+            sample_worksheet = workbook["样本数据"]
+            self.assertEqual(sample_worksheet["B1"].value, "00123")
+            self.assertEqual(sample_worksheet["B1"].data_type, "s")
+
+    def test_survival_chart_uses_pipe_age_x_axis_label(self) -> None:
+        df = pd.DataFrame({ID_COLUMN: ["P001"], PIPE_AGE_COLUMN: [12]})
+
+        with TemporaryDirectory() as temp_dir:
+            output_path = Path(temp_dir) / "chart.png"
+            with patch("app.prediction.plt.xlabel") as xlabel:
+                render_survival_chart(df, [DummyCurve([1, 2], [0.9, 0.7])], output_path)
+
+        xlabel.assert_called_once()
+        self.assertEqual(xlabel.call_args.args[0], "管龄（年）")
 
 
 if __name__ == "__main__":