Spaces:
Running
Running
| """ | |
| 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.""" | |
| 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" | |