aki_prj23_transparenzregister/tests/ai/ner_service_test.py

"""Tests for checking NER Services."""

from aki_prj23_transparenzregister.ai.ner_service import NerAnalysisService


def test_ner_spacy() -> None:
    """Mock TestNerService."""
    # Create instance of NerAnalysisService with use_spacy=True
    ner_service = NerAnalysisService(
        use_spacy=True, use_transformer=False, use_companylist=False
    )
    # 1st testing
    doc = {"title": "Siemens ist ein Unternehmen."}
    result = ner_service.ner_spacy(doc, ent_type="ORG", doc_attrib="title")
    assert result == {"Siemens": 1}

    # 2nd testing
    doc = {"text": "BASF ist ein großes Unternehmen."}
    result = ner_service.ner_spacy(doc, ent_type="ORG", doc_attrib="text")
    assert result == {"BASF": 1}


def test_ner_company_list() -> None:
    """Mock test_ner_company."""
    # Create instance of NerAnalysisService with use_companylist=True
    ner_service = NerAnalysisService(
        use_spacy=False, use_transformer=False, use_companylist=True
    )

    doc = {"title": "Siemens ist ein Unternehmen."}
    result = ner_service.ner_company_list(doc, ent_type="ORG", doc_attrib="title")
    assert result == {"siemens": 1}

    # 2nd testing
    doc = {"text": "BASF ist ein großes Unternehmen."}
    result = ner_service.ner_company_list(doc, ent_type="ORG", doc_attrib="text")
    assert result == {"basf": 1}


def test_ner_transformer() -> None:
    """Mock test_ner_company."""
    # Create instance of NerAnalysisService with use_use_companylist=True
    ner_service = NerAnalysisService(
        use_spacy=False, use_transformer=True, use_companylist=False
    )

    doc = {"title": "Siemens ist ein Unternehmen."}
    result = ner_service.ner_transformer(doc, ent_type="ORG", doc_attrib="title")
    assert result == {"Siemens": 1}

    # 2nd testing
    doc = {"text": "BASF ist ein großes Unternehmen."}
    result = ner_service.ner_transformer(doc, ent_type="ORG", doc_attrib="text")
    assert result == {"BASF": 1}