Spaces:
Running
Running
File size: 7,552 Bytes
79b0bef 4dc0836 79b0bef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """
tests/test_ner.py
────────────────────────────────────────────────────────────────
Unit tests for src/nlp/ner.py
Covers: Entity dataclass, label normalisation, SpacyNERPipeline
interface (mocked model), factory function, edge cases.
The scispaCy model is never loaded — spacy.load() is mocked so
the suite runs in seconds without the 500 MB model file.
────────────────────────────────────────────────────────────────
"""
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from src.nlp.ner import (
BaseNERPipeline,
Entity,
HybridNERPipeline,
SpacyNERPipeline,
_normalise_label,
build_ner_pipeline,
)
class TestEntity:
"""Entity dataclass serialisation."""
def test_to_dict_all_fields(self):
ent = Entity(text="hypertension", label="DISEASE",
start=10, end=22, confidence=0.95, note_id=7)
d = ent.to_dict()
assert d["text"] == "hypertension"
assert d["label"] == "DISEASE"
assert d["start"] == 10
assert d["end"] == 22
assert d["confidence"] == 0.95
assert d["note_id"] == 7
def test_to_dict_none_confidence(self):
ent = Entity(text="pain", label="SYMPTOM", start=0, end=4)
assert ent.to_dict()["confidence"] is None
def test_to_dict_rounds_confidence(self):
ent = Entity(text="aspirin", label="MEDICATION",
start=0, end=7, confidence=0.94567)
assert ent.to_dict()["confidence"] == 0.946
def test_note_id_defaults_to_none(self):
ent = Entity(text="fever", label="SYMPTOM", start=0, end=5)
assert ent.note_id is None
class TestNormaliseLabel:
"""_normalise_label() mapping logic."""
@pytest.mark.parametrize("raw,expected", [
("DISEASE", "DISEASE"),
("CHEMICAL", "MEDICATION"),
("DRUG", "MEDICATION"),
("ANATOMY", "ANATOMY"),
("ORGAN", "ANATOMY"),
("PROCEDURE", "PROCEDURE"),
])
def test_direct_label_map(self, raw, expected):
assert _normalise_label(raw, "entity") == expected
def test_procedure_keyword_match(self):
assert _normalise_label("UNKNOWN", "knee surgery") == "PROCEDURE"
assert _normalise_label("UNKNOWN", "biopsy of liver") == "PROCEDURE"
assert _normalise_label("UNKNOWN", "CT scan") == "PROCEDURE"
def test_anatomy_keyword_match(self):
assert _normalise_label("UNKNOWN", "left ventricle") == "ANATOMY"
assert _normalise_label("UNKNOWN", "right kidney") == "ANATOMY"
def test_default_to_symptom(self):
assert _normalise_label("UNKNOWN", "fatigue") == "SYMPTOM"
assert _normalise_label("XYZ", "nausea") == "SYMPTOM"
def test_case_insensitive_raw_label(self):
# "pain" is deliberately excluded here -- it's in the SYMPTOM
# fast-path regardless of raw label (bc5cdr over-tags things like
# "knee pain"/"cough" as DISEASE; gold-standard review confirmed
# the fast-path must win). Use a term with no fast-path/keyword
# collision so this only tests raw-label case-insensitivity.
assert _normalise_label("disease", "hypertension") == "DISEASE"
assert _normalise_label("Chemical", "aspirin") == "MEDICATION"
def test_procedure_before_anatomy(self):
# "heart surgery" — surgery keyword wins over heart (anatomy)
assert _normalise_label("UNKNOWN", "heart surgery") == "PROCEDURE"
def _mock_spacy(ents):
"""Return a minimal spaCy model mock producing the given entities."""
mock_doc = MagicMock()
mock_ents = []
for text, start, end, label in ents:
e = MagicMock()
e.text = text
e.start_char = start
e.end_char = end
e.label_ = label
mock_ents.append(e)
mock_doc.ents = mock_ents
mock_nlp = MagicMock()
mock_nlp.return_value = mock_doc
mock_nlp.pipe = MagicMock(return_value=iter([mock_doc]))
return mock_nlp
class TestSpacyNERPipeline:
"""SpacyNERPipeline with mocked model."""
def test_extract_returns_entities(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = _mock_spacy([
("hypertension", 10, 22, "DISEASE"),
("metformin", 30, 39, "CHEMICAL"),
])
ents = pipeline.extract("Patient has hypertension. Takes metformin.")
assert len(ents) == 2
assert ents[0].label == "DISEASE"
assert ents[1].label == "MEDICATION"
def test_extract_sorts_by_start(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = _mock_spacy([
("diabetes", 30, 38, "DISEASE"),
("aspirin", 5, 12, "CHEMICAL"),
])
ents = pipeline.extract("Takes aspirin. Has diabetes.")
assert ents[0].start < ents[1].start
def test_extract_skips_short_tokens(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = _mock_spacy([
("BP", 0, 2, "DISEASE"),
("hypertension", 10, 22, "DISEASE"),
])
ents = pipeline.extract("BP 140. Hypertension noted.")
assert len(ents) == 1
assert ents[0].text == "hypertension"
def test_extract_empty_text(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = MagicMock()
assert pipeline.extract("") == []
assert pipeline.extract(" ") == []
pipeline._nlp.assert_not_called()
def test_extract_batch_uses_pipe(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = _mock_spacy([("hypertension", 0, 12, "DISEASE")])
pipeline.extract_batch(["Note one."], batch_size=8)
pipeline._nlp.pipe.assert_called_once()
def test_extract_batch_empty_input(self):
pipeline = SpacyNERPipeline()
pipeline._nlp = MagicMock()
assert pipeline.extract_batch([]) == []
def test_load_error_raises_oserror(self):
pipeline = SpacyNERPipeline(model_name="nonexistent")
with (
patch("spacy.load", side_effect=OSError("not found")),
pytest.raises(OSError, match="not installed"),
):
pipeline._load_model()
def test_is_subclass_of_base(self):
assert issubclass(SpacyNERPipeline, BaseNERPipeline)
class TestBuildNERPipeline:
"""build_ner_pipeline() factory."""
def test_returns_base_pipeline(self):
assert isinstance(build_ner_pipeline(model_name="en_core_sci_lg"),
BaseNERPipeline)
def test_uses_config_default(self):
# Default is "hybrid" (HybridNERPipeline has no single _model_name --
# it wraps two SpacyNERPipeline instances under _fine/_broad).
from src.utils.config import ModelConfig
pipeline = build_ner_pipeline()
if ModelConfig.ner_model == "hybrid":
assert isinstance(pipeline, HybridNERPipeline)
assert pipeline.model_name.startswith("hybrid(")
else:
assert pipeline._model_name == ModelConfig.ner_model
def test_custom_model_name(self):
pipeline = build_ner_pipeline(model_name="en_ner_bc5cdr_md")
assert pipeline._model_name == "en_ner_bc5cdr_md"
|