import pytest
from app.core.document_handlers.ner_processor import NerProcessor

def test_generate_masked_mapping():
    processor = NerProcessor()
    unique_entities = [
        {'text': '李雷', 'type': '人名'},
        {'text': '李明', 'type': '人名'},
        {'text': '王强', 'type': '人名'},
        {'text': 'Acme Manufacturing Inc.', 'type': '英文公司名', 'industry': 'manufacturing'},
        {'text': 'Google LLC', 'type': '英文公司名'},
        {'text': 'A公司', 'type': '公司名称'},
        {'text': 'B公司', 'type': '公司名称'},
        {'text': 'John Smith', 'type': '英文人名'},
        {'text': 'Elizabeth Windsor', 'type': '英文人名'},
        {'text': '华梦龙光伏项目', 'type': '项目名'},
        {'text': '案号12345', 'type': '案号'},
        {'text': '310101198802080000', 'type': '身份证号'},
        {'text': '9133021276453538XT', 'type': '社会信用代码'},
    ]
    linkage = {
        'entity_groups': [
            {
                'group_id': 'g1',
                'group_type': '公司名称',
                'entities': [
                    {'text': 'A公司', 'type': '公司名称', 'is_primary': True},
                    {'text': 'B公司', 'type': '公司名称', 'is_primary': False},
                ]
            },
            {
                'group_id': 'g2',
                'group_type': '人名',
                'entities': [
                    {'text': '李雷', 'type': '人名', 'is_primary': True},
                    {'text': '李明', 'type': '人名', 'is_primary': False},
                ]
            }
        ]
    }
    mapping = processor._generate_masked_mapping(unique_entities, linkage)
    # 人名
    assert mapping['李雷'].startswith('李某')
    assert mapping['李明'].startswith('李某')
    assert mapping['王强'].startswith('王某')
    # 英文公司名
    assert mapping['Acme Manufacturing Inc.'] == 'MANUFACTURING'
    assert mapping['Google LLC'] == 'COMPANY'
    # 公司名同组
    assert mapping['A公司'] == mapping['B公司']
    assert mapping['A公司'].endswith('公司')
    # 英文人名
    assert mapping['John Smith'] == 'J*** S***'
    assert mapping['Elizabeth Windsor'] == 'E*** W***'
    # 项目名
    assert mapping['华梦龙光伏项目'].endswith('项目')
    # 案号
    assert mapping['案号12345'] == '***'
    # 身份证号
    assert mapping['310101198802080000'] == 'XXXXXX'
    # 社会信用代码
    assert mapping['9133021276453538XT'] == 'XXXXXXXX'