Livrare LOT 1 - Didi

ai_platform/modules/extractors/tests/test_ner.py

@@ -0,0 +1,58 @@
+"""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ă"