58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
"""NER tests with a fake GLiNER predictor (no torch/model download)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from fastapi.testclient import TestClient
|
||
|
||
from extractors.app import app
|
||
from extractors.features import ner
|
||
|
||
|
||
class FakePredictor:
|
||
"""Mimics GLiNER.predict_entities with canned spans."""
|
||
|
||
def __init__(self, ents):
|
||
self._ents = ents
|
||
|
||
def predict_entities(self, text, labels, threshold=0.5):
|
||
return self._ents
|
||
|
||
|
||
_SAMPLE = [
|
||
{"text": "Klaus Iohannis", "label": "persoană", "start": 0, "end": 14, "score": 0.97},
|
||
{"text": "Guvernul României", "label": "instituție publică", "start": 20,
|
||
"end": 37, "score": 0.91},
|
||
]
|
||
|
||
|
||
def test_ner_parses_and_sorts():
|
||
res = ner.analyze("Klaus Iohannis și Guvernul României.", model=FakePredictor(_SAMPLE))
|
||
assert res.ok is True
|
||
assert res.results["count"] == 2
|
||
assert res.results["entities"][0]["text"] == "Klaus Iohannis"
|
||
assert res.results["by_label"]["persoană"] == 1
|
||
assert "persoană×1" in res.evidence[0]
|
||
|
||
|
||
def test_ner_empty_entities():
|
||
res = ner.analyze("text neutru", model=FakePredictor([]))
|
||
assert res.ok is True
|
||
assert res.results["count"] == 0
|
||
assert "No entities" in res.evidence[0]
|
||
|
||
|
||
def test_ner_endpoint_503_without_ml_extra():
|
||
# gliner/torch are not installed in the fast (dev) environment.
|
||
client = TestClient(app)
|
||
r = client.post("/v1/ner", json={"text": "Klaus Iohannis"})
|
||
assert r.status_code == 503
|
||
|
||
|
||
def test_ner_endpoint_with_monkeypatched_model(monkeypatch):
|
||
monkeypatch.setattr(ner, "get_model", lambda: FakePredictor(_SAMPLE))
|
||
client = TestClient(app)
|
||
r = client.post("/v1/ner", json={"text": "Klaus Iohannis și Guvernul României."})
|
||
assert r.status_code == 200
|
||
body = r.json()
|
||
assert body["results"]["count"] == 2
|
||
assert body["results"]["entities"][1]["label"] == "instituție publică"
|