nerserver / tests /test_extract.py
Robin
Claude Sonnet 4.6
feat: smarter fallback โ€” heuristic sufficiency + result merging (v3.1)
2288fd7
Raw History Blame Contribute Delete
23.3 kB
"""
Unit tests โ€” no real model loaded (GLiNER/torch stubbed via conftest.py).
Covers:
- API contract (health, validation, threshold, language field)
- Dual-model routing (ZH โ†’ BERT backend, EN/AR/mixed โ†’ GLiNER backend)
- Optional labels + bilingual auto-expansion
- English / Chinese / Arabic / mixed-language scenarios
- labels_used echo in response
"""
from unittest.mock import MagicMock, patch, PropertyMock
import pytest
from fastapi.testclient import TestClient
from app.main import app
from app.models import Entity
from app.labels import DEFAULT_LABELS, expand_bilingual, labels_to_bert_types
from app.ner import detect_language
# โ”€โ”€ Helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def _ents(*args: Entity) -> tuple[list[Entity], list[str]]:
"""Wrap entities in (entities, labels_used) tuple expected by NERService."""
return list(args), [e.label for e in args]
# โ”€โ”€ Fixture โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
@pytest.fixture()
def client():
"""
Patch NERService so no model is actually loaded.
mock_ner.extract() returns ([], []) by default.
"""
mock_ner = MagicMock()
mock_ner.extract.return_value = ([], [])
with pytest.MonkeyPatch().context() as mp:
mp.setattr("app.main.NERService", lambda *_: mock_ner)
with TestClient(app) as c:
yield c, mock_ner
# โ”€โ”€ System / API contract โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_health(client):
c, _ = client
assert c.get("/api/v1/health").json() == {"status": "ok"}
def test_extract_empty_text_returns_empty(client):
c, _ = client
resp = c.post("/api/v1/extract", json={"text": "", "labels": ["person"]})
assert resp.status_code == 200
assert resp.json()["entities"] == []
def test_extract_labels_optional(client):
"""labels ๅญ—ๆฎตๅฎŒๅ…จไธไผ ๅบ”ๆญฃๅธธ่ฟ”ๅ›ž 200ใ€‚"""
c, mock_ner = client
mock_ner.extract.return_value = ([], DEFAULT_LABELS)
resp = c.post("/api/v1/extract", json={"text": "Some text."})
assert resp.status_code == 200
assert len(resp.json()["labels_used"]) > 0
def test_extract_empty_labels_uses_defaults(client):
"""labels=[] ๆ—ถๅบ”ไฝฟ็”จ้ป˜่ฎคๅŒ่ฏญๆ ‡็ญพ้›†ใ€‚"""
c, mock_ner = client
mock_ner.extract.return_value = ([], DEFAULT_LABELS)
resp = c.post("/api/v1/extract", json={"text": "Hello world.", "labels": []})
assert resp.status_code == 200
assert resp.json()["labels_used"] == DEFAULT_LABELS
def test_extract_threshold_forwarded(client):
c, mock_ner = client
c.post("/api/v1/extract",
json={"text": "Hello", "labels": ["person"], "threshold": 0.8})
mock_ner.extract.assert_called_once_with(
"Hello", ["person"], 0.8, language="auto", min_entities=None
)
def test_extract_invalid_threshold(client):
c, _ = client
assert c.post("/api/v1/extract",
json={"text": "x", "threshold": 1.5}).status_code == 422
def test_extract_language_field_forwarded(client):
c, mock_ner = client
c.post("/api/v1/extract",
json={"text": "ๅŒ—ไบฌๅๅ’ŒๅŒป้™ข", "labels": ["ๅŒป้™ขๅ็งฐ"], "language": "zh"})
mock_ner.extract.assert_called_once_with(
"ๅŒ—ไบฌๅๅ’ŒๅŒป้™ข", ["ๅŒป้™ขๅ็งฐ"], 0.4, language="zh", min_entities=None
)
def test_extract_min_entities_forwarded(client):
c, mock_ner = client
c.post("/api/v1/extract",
json={"text": "้ฉฌไบ‘ๅœจๆญๅทžใ€‚", "language": "zh", "min_entities": 5})
mock_ner.extract.assert_called_once_with(
"้ฉฌไบ‘ๅœจๆญๅทžใ€‚", [], 0.4, language="zh", min_entities=5
)
def test_extract_negative_min_entities_rejected(client):
c, _ = client
resp = c.post("/api/v1/extract",
json={"text": "x", "min_entities": -1})
assert resp.status_code == 422
def test_extract_invalid_language(client):
c, _ = client
assert c.post("/api/v1/extract",
json={"text": "x", "language": "jp"}).status_code == 422
def test_labels_used_echoed(client):
c, mock_ner = client
used = ["ไบบๅๆˆ–ๅง“ๅ", "ๅœฐๅๆˆ–ๅŸŽๅธ‚"]
mock_ner.extract.return_value = ([], used)
resp = c.post("/api/v1/extract", json={"text": "้ฉฌไบ‘ๅœจๆญๅทžใ€‚", "labels": ["ไบบๅๆˆ–ๅง“ๅ"]})
assert resp.json()["labels_used"] == used
def test_entity_response_fields(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="Apple", label="organization", score=0.95, start=0, end=5)
)
resp = c.post("/api/v1/extract",
json={"text": "Apple is great.", "labels": ["organization"]})
e = resp.json()["entities"][0]
assert {"text", "label", "score", "start", "end"} <= e.keys()
assert 0.0 <= e["score"] <= 1.0
assert e["start"] < e["end"]
# โ”€โ”€ Language detection โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_detect_language_english():
assert detect_language("Elon Musk founded SpaceX in California.") == "en"
def test_detect_language_chinese():
assert detect_language("้˜ฟ้‡Œๅทดๅทด้›†ๅ›ขๅˆ›ๅง‹ไบบ้ฉฌไบ‘ไบŽๆญๅทžๅธไปปใ€‚") == "zh"
def test_detect_language_arabic():
assert detect_language("ุฃุนู„ู† ุงู„ุฑุฆูŠุณ ู…ุญู…ุฏ ุจู† ุณู„ู…ุงู† ุนู† ู…ุดุฑูˆุน ู†ูŠูˆู….") == "ar"
def test_detect_language_mixed():
assert detect_language("ๅผ ไผŸๅŠ ๅ…ฅไบ† Google ๅŒ—ไบฌ็ ”ๅ‘ไธญๅฟƒ๏ผŒ่ดŸ่ดฃ Android ไผ˜ๅŒ–ใ€‚") == "mixed"
def test_detect_language_empty():
assert detect_language("") == "en"
# โ”€โ”€ Bilingual label expansion โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_expand_bilingual_adds_english_for_chinese():
result = expand_bilingual(["ไบบๅๆˆ–ๅง“ๅ"])
assert "full name of a person" in result
def test_expand_bilingual_adds_chinese_for_english():
result = expand_bilingual(["company or organization name"])
assert "ๅ…ฌๅธๆˆ–็ป„็ป‡ๆœบๆž„ๅ็งฐ" in result
def test_expand_bilingual_no_duplicate():
result = expand_bilingual(["ไบบๅๆˆ–ๅง“ๅ", "full name of a person"])
assert result.count("ไบบๅๆˆ–ๅง“ๅ") == 1
assert result.count("full name of a person") == 1
def test_expand_bilingual_custom_label_preserved():
result = expand_bilingual(["my custom label"])
assert "my custom label" in result
def test_default_labels_bilingual():
has_en = any(all(ord(c) < 128 for c in lbl) for lbl in DEFAULT_LABELS)
has_zh = any(any('ไธ€' <= c <= '้ฟฟ' for c in lbl) for lbl in DEFAULT_LABELS)
assert has_en and has_zh
# โ”€โ”€ BERT label mapping โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_labels_to_bert_types_chinese_label():
types = labels_to_bert_types(["ไบบๅๆˆ–ๅง“ๅ"])
assert "PER" in types
def test_labels_to_bert_types_english_label():
types = labels_to_bert_types(["geographical location"])
assert "LOC" in types
def test_labels_to_bert_types_empty_returns_none():
assert labels_to_bert_types([]) is None
def test_labels_to_bert_types_unmapped_returns_none():
# ๆ— ๆณ•ๆ˜ ๅฐ„็š„ๆ ‡็ญพ โ†’ ไธ่ฟ‡ๆปค๏ผˆ่ฟ”ๅ›ž None๏ผ‰
assert labels_to_bert_types(["some unknown label"]) is None
# โ”€โ”€ English โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_english_person_org(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="Elon Musk", label="person", score=0.98, start=0, end=9),
Entity(text="Tesla", label="organization", score=0.96, start=18, end=23),
Entity(text="SpaceX", label="organization", score=0.97, start=28, end=34),
)
resp = c.post("/api/v1/extract",
json={"text": "Elon Musk is the CEO of Tesla and founded SpaceX.",
"labels": ["full name of a person", "company or organization name"],
"language": "en"})
assert resp.status_code == 200
texts = {e["text"] for e in resp.json()["entities"]}
assert {"Elon Musk", "Tesla", "SpaceX"} <= texts
def test_english_location_date(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="Paris", label="location", score=0.94, start=20, end=25),
Entity(text="2024", label="date", score=0.91, start=29, end=33),
Entity(text="France", label="location", score=0.93, start=38, end=44),
)
resp = c.post("/api/v1/extract",
json={"text": "The summit was held in Paris in 2024, in France.",
"labels": ["geographical location", "date or year"],
"language": "en"})
texts = {e["text"] for e in resp.json()["entities"]}
assert {"Paris", "France", "2024"} <= texts
def test_english_threshold_forwarded(client):
c, mock_ner = client
c.post("/api/v1/extract",
json={"text": "NASA explored the Moon.",
"labels": ["company or organization name"],
"threshold": 0.8, "language": "en"})
mock_ner.extract.assert_called_once_with(
"NASA explored the Moon.", ["company or organization name"], 0.8,
language="en", min_entities=None,
)
# โ”€โ”€ Chinese (BERT backend) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_chinese_person_org(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="ไบบๅๆˆ–ๅง“ๅ", score=0.96, start=8, end=10),
Entity(text="ๅผ ๅ‹‡", label="ไบบๅๆˆ–ๅง“ๅ", score=0.94, start=25, end=27),
Entity(text="้˜ฟ้‡Œๅทดๅทด", label="ๅ…ฌๅธๆˆ–็ป„็ป‡ๆœบๆž„ๅ็งฐ", score=0.97, start=0, end=4),
)
resp = c.post("/api/v1/extract",
json={"text": "้˜ฟ้‡Œๅทดๅทด้›†ๅ›ขๅˆ›ๅง‹ไบบ้ฉฌไบ‘ๅธไปป๏ผŒ็”ฑๅผ ๅ‹‡ๆŽฅไปปใ€‚",
"labels": ["ไบบๅๆˆ–ๅง“ๅ", "ๅ…ฌๅธๆˆ–็ป„็ป‡ๆœบๆž„ๅ็งฐ"],
"language": "zh"})
texts = {e["text"] for e in resp.json()["entities"]}
assert {"้ฉฌไบ‘", "ๅผ ๅ‹‡", "้˜ฟ้‡Œๅทดๅทด"} <= texts
def test_chinese_entity_boundary(client):
"""BERT NER ๅบ”็ฒพ็กฎๆˆชๆ–ญๅฎžไฝ“่พน็•Œ๏ผŒไธๅซๅŠจ่ฏใ€‚"""
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="ๅฐคๆฐ", label="ไบบๅๆˆ–ๅง“ๅ", score=0.82, start=0, end=2),
Entity(text="็Ž‹็†™ๅ‡ค", label="ไบบๅๆˆ–ๅง“ๅ", score=0.95, start=8, end=11),
)
resp = c.post("/api/v1/extract",
json={"text": "ๅฐคๆฐๆฅ่ฏท๏ผŒ็Ž‹็†™ๅ‡ค็ฌ‘้“๏ผš'ไฝ ๆฅไบ†ใ€‚'",
"labels": ["ไบบๅๆˆ–ๅง“ๅ"], "language": "zh"})
texts = {e["text"] for e in resp.json()["entities"]}
assert "ๅฐคๆฐ" in texts
assert "็Ž‹็†™ๅ‡ค" in texts
assert "ๅฐคๆฐๆฅ่ฏท" not in texts
assert "็Ž‹็†™ๅ‡ค็ฌ‘้“" not in texts
def test_chinese_location_product(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="ๆญๅทž", label="ๅœฐๅๆˆ–ๅŸŽๅธ‚", score=0.93, start=9, end=11),
Entity(text="ๆท˜ๅฎ", label="ไบงๅ“ๆˆ–ๅ“็‰Œๅ็งฐ", score=0.91, start=14, end=16),
Entity(text="ๅคฉ็Œซ", label="ไบงๅ“ๆˆ–ๅ“็‰Œๅ็งฐ", score=0.92, start=17, end=19),
Entity(text="ๆ”ฏไป˜ๅฎ", label="ไบงๅ“ๆˆ–ๅ“็‰Œๅ็งฐ", score=0.90, start=20, end=23),
)
resp = c.post("/api/v1/extract",
json={"text": "้˜ฟ้‡Œๅทดๅทดๆ€ป้ƒจไฝไบŽๆญๅทž๏ผŒๆ——ไธ‹ๆœ‰ๆท˜ๅฎใ€ๅคฉ็Œซใ€ๆ”ฏไป˜ๅฎใ€‚",
"labels": ["ๅœฐๅๆˆ–ๅŸŽๅธ‚", "ไบงๅ“ๆˆ–ๅ“็‰Œๅ็งฐ"],
"language": "zh"})
texts = {e["text"] for e in resp.json()["entities"]}
assert {"ๆญๅทž", "ๆท˜ๅฎ", "ๅคฉ็Œซ", "ๆ”ฏไป˜ๅฎ"} <= texts
def test_chinese_auto_routes_to_zh(client):
"""auto ๆฃ€ๆต‹ๅˆฐไธญๆ–‡ๅบ”่ทฏ็”ฑๅˆฐ ZH ๆจกๅž‹๏ผˆlanguage ้€ไผ ไธบ 'auto'๏ผŒๅ†…้ƒจๆฃ€ๆต‹ไธบ zh๏ผ‰ใ€‚"""
c, mock_ner = client
c.post("/api/v1/extract",
json={"text": "้ฉฌไบ‘ๅˆ›็ซ‹ไบ†้˜ฟ้‡Œๅทดๅทดใ€‚"})
# NERService.extract ่ขซ่ฐƒ็”จๆ—ถ language='auto'๏ผŒ่ทฏ็”ฑ้€ป่พ‘ๅœจ ner.py ๅ†…้ƒจๅค„็†
mock_ner.extract.assert_called_once()
call_kwargs = mock_ner.extract.call_args
assert call_kwargs[1].get("language") == "auto" or call_kwargs[0][3] == "auto"
# โ”€โ”€ Arabic โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_arabic_person_location(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="ู…ุญู…ุฏ ุจู† ุณู„ู…ุงู†", label="full name of a person", score=0.82, start=12, end=26),
Entity(text="ุงู„ู…ู…ู„ูƒุฉ ุงู„ุนุฑุจูŠุฉ ุงู„ุณุนูˆุฏูŠุฉ", label="geographical location", score=0.85, start=44, end=68),
)
resp = c.post("/api/v1/extract",
json={"text": "ุฃุนู„ู† ุงู„ุฑุฆูŠุณ ู…ุญู…ุฏ ุจู† ุณู„ู…ุงู† ู…ุดุฑูˆุน ู†ูŠูˆู… ููŠ ุงู„ู…ู…ู„ูƒุฉ ุงู„ุนุฑุจูŠุฉ ุงู„ุณุนูˆุฏูŠุฉ.",
"labels": ["full name of a person", "geographical location"],
"language": "ar"})
texts = {e["text"] for e in resp.json()["entities"]}
assert "ู…ุญู…ุฏ ุจู† ุณู„ู…ุงู†" in texts
assert "ุงู„ู…ู…ู„ูƒุฉ ุงู„ุนุฑุจูŠุฉ ุงู„ุณุนูˆุฏูŠุฉ" in texts
# โ”€โ”€ Mixed Chinese-English โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_mixed_entities_both_scripts(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="ๅผ ไผŸ", label="person", score=0.95, start=0, end=2),
Entity(text="Google", label="organization", score=0.97, start=9, end=15),
Entity(text="ๅŒ—ไบฌ", label="location", score=0.93, start=25, end=27),
Entity(text="Android", label="product", score=0.91, start=33, end=40),
)
resp = c.post("/api/v1/extract",
json={"text": "ๅผ ไผŸๅ…ฅ่Œไบ† Google๏ผŒ้ฉปๆ‰ŽๅœจๅŒ—ไบฌ๏ผŒ่ดŸ่ดฃ Android ็ ”ๅ‘ใ€‚",
"labels": ["full name of a person", "ไบบๅๆˆ–ๅง“ๅ",
"company or organization name", "ๅ…ฌๅธๆˆ–็ป„็ป‡ๆœบๆž„ๅ็งฐ",
"geographical location", "ๅœฐๅๆˆ–ๅŸŽๅธ‚",
"product or technology name"],
"language": "mixed"})
texts = {e["text"] for e in resp.json()["entities"]}
assert {"ๅผ ไผŸ", "Google", "ๅŒ—ไบฌ", "Android"} <= texts
def test_mixed_no_cross_language_contamination(client):
c, mock_ner = client
mock_ner.extract.return_value = _ents(
Entity(text="OpenAI", label="organization", score=0.97, start=5, end=11),
Entity(text="็Ž‹่Šณ", label="person", score=0.93, start=15, end=17),
)
resp = c.post("/api/v1/extract",
json={"text": "ไป–ๅœจ OpenAI ๅทฅไฝœ๏ผŒๅŒไบ‹็Ž‹่ŠณไนŸๅœจๅŒไธ€้ƒจ้—จใ€‚",
"labels": ["person", "organization"]})
entities = resp.json()["entities"]
assert any(e["text"] == "OpenAI" and e["label"] == "organization" for e in entities)
assert any(e["text"] == "็Ž‹่Šณ" and e["label"] == "person" for e in entities)
# โ”€โ”€ Fallback & merge (NERService unit tests, no HTTP) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def _build_svc():
"""Construct a bare NERService with mocked backends and locks."""
import threading
from app.ner import NERService
svc = NERService.__new__(NERService)
svc._en_lock = threading.Lock()
svc._zh_lock = threading.Lock()
svc._en_backend = MagicMock()
svc._zh_backend = MagicMock()
return svc
# โ”€โ”€ Sufficiency heuristic โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_expected_min_short_text():
from app.ner import NERService
assert NERService._expected_min("้ฉฌไบ‘", []) == 1
def test_expected_min_medium_text():
from app.ner import NERService
text = "x" * 50
assert NERService._expected_min(text, []) == 2
def test_expected_min_long_text():
from app.ner import NERService
text = "x" * 350
assert NERService._expected_min(text, []) == 4
def test_expected_min_label_floor_takes_over():
"""9 ไธชๆ ‡็ญพ โ†’ โŒˆ9/3โŒ‰=3๏ผŒ่ถ…่ฟ‡็Ÿญๆ–‡ๆœฌ็š„ length_floor=1๏ผŒๆœ€็ปˆๅ– 3ใ€‚"""
from app.ner import NERService
short_text = "้ฉฌไบ‘"
labels = [f"l{i}" for i in range(9)]
assert NERService._expected_min(short_text, labels) == 3
# โ”€โ”€ ZH branch fallback โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_zh_empty_triggers_fallback_and_adds():
"""ZH ไธปๆจกๅž‹ 0 ไธช โ†’ ่งฆๅ‘ๅ…œๅบ• โ†’ ่ฟ”ๅ›žๅ…œๅบ•็ป“ๆžœใ€‚"""
svc = _build_svc()
svc._zh_backend.predict.return_value = ([], [])
svc._en_backend.predict.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="person", score=0.75, start=0, end=2)
)
entities, _ = svc.extract("้ฉฌไบ‘", [], 0.4, language="zh")
assert any(e.text == "้ฉฌไบ‘" for e in entities)
svc._zh_backend.predict.assert_called_once()
svc._en_backend.predict.assert_called_once()
def test_zh_sufficient_no_fallback():
"""ZH ไธปๆจกๅž‹ๅฎžไฝ“ๆ•ฐ โ‰ฅ expected_min(=1 ็Ÿญๆ–‡ๆœฌ) โ†’ ไธ่ฐƒ็”จๅ…œๅบ•ใ€‚"""
svc = _build_svc()
svc._zh_backend.predict.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="person", score=0.92, start=0, end=2)
)
svc.extract("้ฉฌไบ‘", [], 0.4, language="zh")
svc._en_backend.predict.assert_not_called()
def test_zh_insufficient_triggers_fallback_and_results_added():
"""
ๅ…ณ้”ฎๆต‹่ฏ•๏ผšZH ่ฟ”ๅ›ž 1 ไธช๏ผŒไฝ†ๆ–‡ๆœฌ้•ฟ โ†’ expected_min=4๏ผŒไธๅ……ๅˆ† โ†’
่งฆๅ‘ๅ…œๅบ•๏ผŒไธป็ป“ๆžœ + ๅ…œๅบ•็ป“ๆžœไธ€ๅนถ่ฟ”ๅ›ž๏ผˆ็›ธๅŠ ๏ผŒไธๆ›ฟๆข๏ผ‰ใ€‚
"""
svc = _build_svc()
long_text = "้ฉฌไบ‘" + "x" * 350 # length_floor = 4
svc._zh_backend.predict.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="ไบบๅๆˆ–ๅง“ๅ", score=0.95, start=0, end=2),
)
svc._en_backend.predict.return_value = _ents(
Entity(text="Tesla", label="organization", score=0.90, start=10, end=15),
Entity(text="2024", label="date", score=0.88, start=20, end=24),
)
entities, _ = svc.extract(long_text, [], 0.4, language="zh")
texts = {e.text for e in entities}
# ไธปๆจกๅž‹็š„"้ฉฌไบ‘"ๅฟ…้กปไฟ็•™๏ผŒๅŒๆ—ถๅ…œๅบ•็š„ Tesla / 2024 ไนŸๅŠ ่ฟ›ๆฅ
assert "้ฉฌไบ‘" in texts
assert "Tesla" in texts
assert "2024" in texts
assert len(entities) == 3 # 1 + 2 = 3๏ผŒ็กฎๅฎžๆ˜ฏ็›ธๅŠ 
def test_user_min_entities_overrides_heuristic():
"""่ฏทๆฑ‚้‡Œไผ  min_entities=5 ๆ—ถๅบ”่ฆ†็›–ๅฏๅ‘ๅผ๏ผŒไธปๆจกๅž‹ 3 ไธชไป่งฆๅ‘ๅ…œๅบ•ใ€‚"""
svc = _build_svc()
svc._zh_backend.predict.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="person", score=0.95, start=0, end=2),
Entity(text="ๅผ ๅ‹‡", label="person", score=0.93, start=4, end=6),
Entity(text="ๆญๅทž", label="location", score=0.91, start=8, end=10),
)
svc._en_backend.predict.return_value = _ents(
Entity(text="Tesla", label="organization", score=0.85, start=15, end=20),
)
entities, _ = svc.extract("้ฉฌไบ‘ใ€ๅผ ๅ‹‡ใ€ๆญๅทž Tesla", [], 0.4,
language="zh", min_entities=5)
# 3 < 5 โ†’ ่งฆๅ‘ๅ…œๅบ•๏ผ›ๆœ€็ปˆ 3 + 1 = 4
assert len(entities) == 4
svc._en_backend.predict.assert_called_once()
def test_user_min_entities_zero_disables_fallback():
"""min_entities=0 ๆ—ถไธปๆจกๅž‹ๅณไฝฟ่ฟ”ๅ›ž็ฉบไนŸไธ่งฆๅ‘ๅ…œๅบ•ใ€‚"""
svc = _build_svc()
svc._zh_backend.predict.return_value = ([], [])
svc.extract("้ฉฌไบ‘", [], 0.4, language="zh", min_entities=0)
svc._en_backend.predict.assert_not_called()
# โ”€โ”€ EN branch fallback (symmetric) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_en_insufficient_triggers_zh_fallback_and_adds():
"""EN ไธปๆจกๅž‹ไธๅ……ๅˆ† โ†’ ่ฐƒ ZH ๅ…œๅบ• โ†’ ็ป“ๆžœ็›ธๅŠ ใ€‚"""
svc = _build_svc()
long_text = "Tesla" + "x" * 350
svc._en_backend.predict.return_value = _ents(
Entity(text="Tesla", label="organization", score=0.95, start=0, end=5),
)
svc._zh_backend.predict.return_value = _ents(
Entity(text="้ฉฌไบ‘", label="ไบบๅๆˆ–ๅง“ๅ", score=0.91, start=10, end=12),
)
entities, _ = svc.extract(long_text, [], 0.4, language="en")
texts = {e.text for e in entities}
assert "Tesla" in texts and "้ฉฌไบ‘" in texts
assert len(entities) == 2
# โ”€โ”€ Mixed: always merge both โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
def test_mixed_always_runs_both_models():
"""mixed ่ฏญ่จ€ๆ— ่ง†ๅ……ๅˆ†ๆ€ง๏ผŒๆฐธ่ฟœ่ท‘ไธคไธชๆจกๅž‹ๅนถๅˆๅนถใ€‚"""
svc = _build_svc()
svc._en_backend.predict.return_value = _ents(
Entity(text="Google", label="organization", score=0.95, start=5, end=11)
)
svc._zh_backend.predict.return_value = _ents(
Entity(text="ๅผ ไผŸ", label="ไบบๅๆˆ–ๅง“ๅ", score=0.91, start=0, end=2)
)
entities, _ = svc.extract("ๅผ ไผŸๅŠ ๅ…ฅ Googleใ€‚", [], 0.4, language="mixed")
texts = {e.text for e in entities}
assert {"Google", "ๅผ ไผŸ"} <= texts
svc._en_backend.predict.assert_called_once()
svc._zh_backend.predict.assert_called_once()
def test_merge_deduplicates_overlapping_spans():
"""ไธคไธชๆจกๅž‹ๅฏนๅŒไธ€ span ้ƒฝๅ‘ฝไธญ โ†’ ไฟ็•™ๅพ—ๅˆ†ๆœ€้ซ˜็š„้‚ฃๆกใ€‚"""
svc = _build_svc()
svc._en_backend.predict.return_value = (
[Entity(text="ๅผ ไผŸ", label="person", score=0.70, start=0, end=2)],
["person"],
)
svc._zh_backend.predict.return_value = (
[Entity(text="ๅผ ไผŸ", label="ไบบๅๆˆ–ๅง“ๅ", score=0.92, start=0, end=2)],
["ไบบๅๆˆ–ๅง“ๅ"],
)
entities, _ = svc.extract("ๅผ ไผŸ", [], 0.4, language="mixed")
matches = [e for e in entities if e.text == "ๅผ ไผŸ"]
assert len(matches) == 1
assert matches[0].score == 0.92 # ้ซ˜ๅˆ†่ƒœๅ‡บ