286 lines
12 KiB
Python
286 lines
12 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 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,
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
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, 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级", "预计剩余寿命（年）": ">63"},
|
|
{ID_COLUMN: "P002", PIPE_AGE_COLUMN: "12 年", "当前健康状态": 0.6, "健康等级": "III级", "预计剩余寿命（年）": ">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, "当前健康状态", "健康等级", "预计剩余寿命（年）"],
|
|
)
|
|
|
|
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:
|
|
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级", "预计剩余寿命（年）": ">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,
|
|
patch("app.prediction.plt.title") as title,
|
|
):
|
|
summary_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级")
|
|
|
|
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,
|
|
}
|
|
)
|
|
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()
|