162 lines
6.5 KiB
Python
162 lines
6.5 KiB
Python
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,
|
|
)
|
|
|
|
|
|
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_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]),
|
|
DummyCurve([1, 2], [0.8, 0.6]),
|
|
]
|
|
summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}]
|
|
summary_sheet_rows = [
|
|
{"管道编号": "P001", "管龄": "10 年", "健康概率": 0.7, "预计剩余寿命": 2.0, "健康等级": "IV级"},
|
|
{"管道编号": "P002", "管龄": "12 年", "健康概率": 0.6, "预计剩余寿命": 1.5, "健康等级": "III级"},
|
|
]
|
|
|
|
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 年", "健康等级": "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_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:
|
|
render_survival_chart(df, [DummyCurve([1, 2], [0.9, 0.7])], output_path)
|
|
|
|
xlabel.assert_called_once()
|
|
self.assertEqual(xlabel.call_args.args[0], "管龄（年）")
|
|
ylabel.assert_called_once()
|
|
self.assertEqual(ylabel.call_args.args[0], "健康风险")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|