lingyue-zhibao/python-services/ner-service/tests/test_ner.py

"""
NER 服务测试
"""
import pytest
from fastapi.testclient import TestClient

from app.main import app


client = TestClient(app)


def test_health_check():
    """测试健康检查接口"""
    response = client.get("/health")
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "ok"
    assert "version" in data


def test_extract_entities():
    """测试实体提取接口"""
    request_data = {
        "documentId": "test-doc-001",
        "text": "2024年5月15日，成都检测公司在成都市高新区完成了环境监测项目的检测工作，使用了噪音检测设备。",
        "extractRelations": True
    }
    
    response = client.post("/ner/extract", json=request_data)
    assert response.status_code == 200
    
    data = response.json()
    assert data["success"] is True
    assert data["documentId"] == "test-doc-001"
    assert "entities" in data
    assert len(data["entities"]) > 0
    
    # 验证提取到的实体类型
    entity_types = {e["type"] for e in data["entities"]}
    assert "DATE" in entity_types or "ORG" in entity_types


def test_extract_relations():
    """测试关系抽取接口"""
    request_data = {
        "documentId": "test-doc-002",
        "text": "成都检测公司负责环境监测项目",
        "entities": [
            {
                "name": "成都检测公司",
                "type": "ORG",
                "value": "成都检测公司",
                "position": {"charStart": 0, "charEnd": 6, "line": 1},
                "tempId": "e1"
            },
            {
                "name": "环境监测项目",
                "type": "PROJECT",
                "value": "环境监测项目",
                "position": {"charStart": 8, "charEnd": 14, "line": 1},
                "tempId": "e2"
            }
        ]
    }
    
    response = client.post("/ner/relations", json=request_data)
    assert response.status_code == 200
    
    data = response.json()
    assert data["success"] is True
    assert "relations" in data


def test_empty_text():
    """测试空文本"""
    request_data = {
        "documentId": "test-doc-003",
        "text": "",
        "extractRelations": False
    }
    
    response = client.post("/ner/extract", json=request_data)
    assert response.status_code == 200
    
    data = response.json()
    assert data["success"] is True
    assert len(data["entities"]) == 0


def test_text_too_long():
    """测试文本过长"""
    request_data = {
        "documentId": "test-doc-004",
        "text": "a" * 60000,  # 超过限制
        "extractRelations": False
    }
    
    response = client.post("/ner/extract", json=request_data)
    assert response.status_code == 400  # 应该返回错误


if __name__ == "__main__":
    pytest.main([__file__, "-v"])