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, 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, ) 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") 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) 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_columns_match_backend_contract(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.assertEqual(len(headers), len(INPUT_COLUMNS)) self.assertEqual(set(headers), set(INPUT_COLUMNS)) 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, 0, 5, 1, 100, 1.2, 0.4, 20, 800, 1, "=缺陷计算!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_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]), DummyCurve([1, 2], [0.8, 0.6]), ] summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}] summary_sheet_rows = [ {ID_COLUMN: "P001", PIPE_AGE_COLUMN: "10 年", "健康风险值": 0.3}, {ID_COLUMN: "P002", PIPE_AGE_COLUMN: "12 年", "健康风险值": 0.4}, ] 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, ["结果摘要", "样本数据"]) 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(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 年", "健康风险值": 0.1}] 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_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, patch("app.prediction.plt.title") as title, ): summary_rows, _ = render_survival_chart(df, [DummyCurve([1, 10, 12], [0.9, 0.7, 0.4])], output_path) self.assertAlmostEqual(summary_rows[0]["health_probability"], 0.4) self.assertAlmostEqual(summary_rows[0]["health_risk"], 0.6) self.assertEqual(summary_rows[0]["grade_label"], "II级") xlabel.assert_called_once() self.assertEqual(xlabel.call_args.args[0], "管龄（年）") ylabel.assert_called_once() self.assertEqual(ylabel.call_args.args[0], "健康风险") title.assert_called_once() self.assertEqual(title.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, } ) curves = [DummyCurve([1, 10, 12], [0.9, 0.7, 0.4]) for _ 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(df, curves, 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()