test_data_normalization_layer.py

#!/usr/bin/env python3
"""Regression tests for unified data normalization layer."""

from __future__ import annotations

from utils.data_normalization import (
    build_risk_aliases,
    canonical_text_key,
    dedupe_keep_order,
    extract_smoking_status_from_text,
    normalize_gender,
    normalize_risk_name,
    normalize_smoking_status,
    normalize_token,
)


def test_normalize_gender_tokens() -> None:
    assert normalize_gender("1") == "male"
    assert normalize_gender("男") == "male"
    assert normalize_gender("female") == "female"
    assert normalize_gender("未知") == "unknown"


def test_normalize_token_and_dedupe() -> None:
    assert normalize_token("  'ABC'  ") == "abc"
    assert dedupe_keep_order([" a ", "a", "", None, "b"]) == ["a", "b"]


def test_smoking_status_normalization() -> None:
    assert normalize_smoking_status("current_smoker") == "current_smoker"
    assert normalize_smoking_status("吸烟状态:不吸烟") == "non_smoker"
    assert normalize_smoking_status("smoking_status:passive_smoker") == "passive_smoker"
    assert normalize_smoking_status("unknown") is None


def test_extract_smoking_status_from_text() -> None:
    text = "患者补问：吸烟状态:不吸烟但有被动吸烟；其他无"
    assert extract_smoking_status_from_text(text) == "passive_smoker"


def test_risk_name_normalization_and_aliases() -> None:
    assert normalize_risk_name("慢阻肺") == "慢性阻塞性肺疾病"
    assert normalize_risk_name("乳腺癌（一般风险人群）") == "乳腺癌"

    aliases = build_risk_aliases(["2型糖尿病", "血管疾病", "心血管病"])
    assert "糖尿病" in aliases
    assert "2型糖尿病" in aliases
    assert "血管疾病筛查" in aliases
    assert "高血压" in aliases


def test_canonical_text_key() -> None:
    assert canonical_text_key("  LDL-C（低密度）: 检查项 ") == canonical_text_key("LDL-C低密度检查项")