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 openpyxl import Workbook from app.prediction import ( CHART_DISPLAY_LIMIT, FEATURES, ID_COLUMN, INPUT_COLUMNS, PIPE_AGE_COLUMN, REMAINING_LIFE_COLUMN, PredictionError, estimate_remaining_life, format_years, grade_info, interpolate_probability, make_analysis_text, prepare_model_features, read_input_file, render_survival_chart, secure_upload_name, validate_input_frame, write_prediction_workbook, ) def dummy_prediction( pipe_id: str = "P001", pipe_age: list[float] | None = None, health_state: list[float] | None = None, current_age: float = 12.0, current_health_state: float = 0.4, ) -> dict: return { "ID": pipe_id, "current_age": current_age, "current_health_state": current_health_state, "current_health_grade": "II级", "pipe_age": pipe_age or [1, 10, 12], "health_state": health_state or [0.9, 0.7, 0.4], } class PredictionHelpersTest(unittest.TestCase): def test_secure_upload_name_accepts_chinese_filename(self) -> None: filename, suffix = secure_upload_name("管道数据.xlsx", "run123") self.assertEqual(suffix, ".xlsx") self.assertTrue(filename.endswith("_run123.xlsx")) def test_secure_upload_name_rejects_unsupported_extension(self) -> None: with self.assertRaises(PredictionError): secure_upload_name("管道数据.txt", "run123") def test_probability_helpers(self) -> None: self.assertEqual(interpolate_probability([1, 5, 10], [0.9, 0.8, 0.6], 6), 0.6) self.assertEqual(estimate_remaining_life([1, 5, 10], [0.9, 0.4, 0.2]), 5.0) self.assertEqual(format_years(None, 10, 73), 63.0) def test_analysis_text_displays_health_state_as_decimal(self) -> None: text = make_analysis_text( [ {"pipe_id": "P001", "health_state": 0.734, "grade_label": "IV级"}, {"pipe_id": "P002", "health_state": 0.2, "grade_label": "I级"}, ] ) self.assertIn("0.73", text) self.assertNotIn("%", text) def test_grade_boundaries(self) -> None: self.assertEqual(grade_info(0.2)[0], "I级") self.assertEqual(grade_info(0.8)[0], "IV级") self.assertEqual(grade_info(0.81)[0], "V级") def test_validate_input_frame_reports_missing_columns(self) -> None: df = pd.DataFrame({ID_COLUMN: [1], FEATURES[0]: [1]}) with self.assertRaises(PredictionError) as ctx: validate_input_frame(df) self.assertIn("缺少必要字段", ctx.exception.message) def test_validate_input_frame_accepts_current_template_columns(self) -> None: df = pd.DataFrame([{column: 1 for column in INPUT_COLUMNS}]) validate_input_frame(df) def test_example_workbook_contains_backend_contract_columns(self) -> None: workbook = load_workbook(Path(__file__).resolve().parents[1] / "example.xlsx", read_only=True) worksheet = workbook["Template"] headers = [cell.value for cell in next(worksheet.iter_rows(min_row=1, max_row=1))] self.assertTrue(set(INPUT_COLUMNS).issubset(headers)) 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_read_input_file_calculates_defects_from_template_detail_sheet(self) -> None: with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "input.xlsx" workbook = Workbook() template = workbook.active template.title = "Template" template.append(INPUT_COLUMNS) template.append(["001", 5, 1, 100, 1, 1.2, 0.4, 20, 800, "=缺陷计算!E2", "=缺陷计算!J2"]) detail = workbook.create_sheet("缺陷计算") detail.append([ ID_COLUMN, "泄漏\n权重0.5", "腐蚀\n权重0.4", "管瘤\n权重0.1", "结构性缺陷值", "气囊\n权重0.5", "杂质\n权重0.2", "异物\n权重0.2", "不明连接\n权重0.1", "功能性缺陷值", ]) detail.append(["001", "严重", "中度", "轻度", None, "轻度", "无", "中度", "严重", None]) workbook.save(output_path) df = read_input_file(output_path, ".xlsx") self.assertAlmostEqual(float(df.loc[0, "结构缺陷"]), 3.8) self.assertAlmostEqual(float(df.loc[0, "功能缺陷"]), 1.6) def test_prepare_model_features_maps_material_aliases_to_rc1_names(self) -> None: rows = [] for value in ["镀锌", "2-钢塑", 11]: row = {feature: 1 for feature in FEATURES} row["管材"] = value row[ID_COLUMN] = "P001" row[PIPE_AGE_COLUMN] = 12 rows.append(row) df = pd.DataFrame(rows) x_test = prepare_model_features(df) self.assertEqual(x_test["Material"].tolist(), ["镀锌", "钢塑", "钢管"]) def test_prepare_model_features_rejects_invalid_material_alias(self) -> None: row = {feature: 1 for feature in FEATURES} row["管材"] = "未知管材" row[ID_COLUMN] = "P001" row[PIPE_AGE_COLUMN] = 12 df = pd.DataFrame([row]) with self.assertRaises(PredictionError) as ctx: prepare_model_features(df) self.assertIn("管材超出RC1支持范围", ctx.exception.message) def test_prepare_model_features_rejects_unsupported_legacy_material_code(self) -> None: row = {feature: 1 for feature in FEATURES} row["管材"] = 13 row[ID_COLUMN] = "P001" row[PIPE_AGE_COLUMN] = 12 df = pd.DataFrame([row]) with self.assertRaises(PredictionError) as ctx: prepare_model_features(df) self.assertIn("管材超出RC1支持范围", ctx.exception.message) def test_prepare_model_features_maps_location_aliases_to_rc1_names(self) -> None: row = {feature: 1 for feature in FEATURES} row["管材"] = 5 row["位置"] = "2-行人道" row[ID_COLUMN] = "P001" row[PIPE_AGE_COLUMN] = 12 df = pd.DataFrame([row]) x_test = prepare_model_features(df) self.assertEqual(x_test["Location"].tolist(), ["行人道"]) def test_prediction_workbook_keeps_sample_data_in_one_sheet(self) -> None: predictions = [ dummy_prediction("P001", [1, 2], [0.9, 0.7], 10, 0.7), dummy_prediction("P002", [1, 2], [0.8, 0.6], 12, 0.6), ] summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}] summary_sheet_rows = [ {ID_COLUMN: "P001", PIPE_AGE_COLUMN: 10, "当前健康状态": 0.7, "健康等级": "IV级", REMAINING_LIFE_COLUMN: 63}, {ID_COLUMN: "P002", PIPE_AGE_COLUMN: 12, "当前健康状态": 0.6, "健康等级": "III级", REMAINING_LIFE_COLUMN: 61}, ] with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "prediction.xlsx" write_prediction_workbook(output_path, predictions, summary_rows, summary_sheet_rows) 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, "当前健康状态", "健康等级", REMAINING_LIFE_COLUMN], ) summary_records = summary_data[summary_data[PIPE_AGE_COLUMN].notna()] self.assertEqual(summary_records[PIPE_AGE_COLUMN].tolist(), [10, 12]) self.assertEqual(summary_records[REMAINING_LIFE_COLUMN].tolist(), [63, 61]) sample_data = pd.read_excel(output_path, sheet_name="样本数据") 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["样本数据"] summary_worksheet = openpyxl_workbook["结果摘要"] self.assertEqual(sample_worksheet.freeze_panes, "A2") self.assertEqual(sample_worksheet.auto_filter.ref, "A1:C3") self.assertEqual(summary_worksheet["C2"].number_format, "0.00") self.assertEqual(sample_worksheet["B2"].number_format, "0.00") self.assertEqual(summary_worksheet["E1"].value, "剩余寿命（年）") self.assertIn("A5:E8", [str(range_) for range_ in summary_worksheet.merged_cells.ranges]) self.assertIn("73减去当前管龄", summary_worksheet["A5"].value) self.assertIn("目前收集到模型数据管龄最大为73年", summary_worksheet["A5"].value) def test_prediction_workbook_writes_pipe_ids_as_excel_text(self) -> None: predictions = [dummy_prediction("00123", [1], [0.9], 10, 0.9)] summary_rows = [{"pipe_id": "00123"}] summary_sheet_rows = [ {ID_COLUMN: "00123", PIPE_AGE_COLUMN: 10, "当前健康状态": 0.9, "健康等级": "V级", REMAINING_LIFE_COLUMN: 63} ] with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "prediction.xlsx" write_prediction_workbook(output_path, predictions, 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_expected_axis_labels(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, patch("app.prediction.plt.ylabel") as ylabel, patch("app.prediction.plt.figtext") as figtext, ): summary_rows, summary_sheet_rows = render_survival_chart([dummy_prediction()], output_path) self.assertAlmostEqual(summary_rows[0]["health_state"], 0.4) self.assertEqual(summary_rows[0]["grade_label"], "II级") self.assertEqual(summary_sheet_rows[0][PIPE_AGE_COLUMN], 12.0) xlabel.assert_called_once() self.assertEqual(xlabel.call_args.args[0], "管龄（年）") ylabel.assert_called_once() self.assertEqual(ylabel.call_args.args[0], "管道健康状态") figtext.assert_called_once() self.assertIn(f"前{CHART_DISPLAY_LIMIT}条管道的示例数据", figtext.call_args.args[2]) def test_survival_chart_only_plots_first_ten_pipes(self) -> None: sample_count = CHART_DISPLAY_LIMIT + 2 df = pd.DataFrame( { ID_COLUMN: [f"P{i:03d}" for i in range(sample_count)], PIPE_AGE_COLUMN: [12] * sample_count, } ) predictions = [ dummy_prediction(f"P{i:03d}", [1, 10, 12], [0.9, 0.7, 0.4]) for i in range(sample_count) ] with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "chart.png" with patch("app.prediction.plt.step") as step, patch("app.prediction.plt.legend"): summary_rows, summary_sheet_rows = render_survival_chart(predictions, output_path) self.assertEqual(step.call_count, CHART_DISPLAY_LIMIT) self.assertEqual(len(summary_rows), sample_count) self.assertEqual(len(summary_sheet_rows), sample_count) if __name__ == "__main__": unittest.main()