tests/test_normalizer.py

import pytest
import uuid
import json
from unittest.mock import AsyncMock, MagicMock, patch

from src.models.claim import Claim, ClaimType, Polarity, StudyDesign, Entity, EntityType
from src.entity.normalizer import EntityNormalizer, EntityNormalizationResponse, NormalizedEntity


@pytest.fixture
def sample_claims():
    entity_asa = Entity(text="ASA", entity_type=EntityType.DRUG)
    entity_acid = Entity(text="acetylsalicylic acid", entity_type=EntityType.DRUG)
    entity_aspirin = Entity(text="aspirin", entity_type=EntityType.DRUG)
    
    entity_metformin = Entity(text="glucophage", entity_type=EntityType.DRUG)
    entity_insulin = Entity(text="insulin", entity_type=EntityType.DRUG) # not in local map, will trigger LLM fallback
    
    claim_1 = Claim(
        id=uuid.uuid4(),
        text="ASA reduces headache.",
        paper_id="11111",
        year=2020,
        confidence_score=1.0,
        claim_type=ClaimType.CAUSAL,
        polarity=Polarity.NEGATIVE,
        entities=[entity_asa, entity_acid],
        population="humans",
        context="general",
        quote_anchor="reduces headache",
        study_design=StudyDesign.RCT
    )
    
    claim_2 = Claim(
        id=uuid.uuid4(),
        text="Aspirin and glucophage treatment.",
        paper_id="22222",
        year=2023,
        confidence_score=1.0,
        claim_type=ClaimType.CAUSAL,
        polarity=Polarity.POSITIVE,
        entities=[entity_aspirin, entity_metformin, entity_insulin],
        population="humans",
        context="general",
        quote_anchor="treatment",
        study_design=StudyDesign.RCT
    )
    
    return [claim_1, claim_2]


@pytest.mark.asyncio
async def test_normalize_entities_local_lookup(sample_claims):
    # Ensure that local lookup works and doesn't hit the LLM since we don't mock it here
    # (except for 'insulin', which we should mock or test separately. So let's temporarily mock the LLM
    # call to return a mock response for 'insulin')
    mock_llm = MagicMock()
    mock_llm.model_name = "mock-llm"
    mock_llm.generate_structured = AsyncMock(
        return_value=EntityNormalizationResponse(
            normalized_entities=[
                NormalizedEntity(
                    original_text="insulin",
                    canonical_name="Insulin",
                    canonical_id="MeSH:D007328",
                    entity_type="DRUG"
                )
            ]
        )
    )
    
    with patch("src.entity.normalizer.get_llm", return_value=mock_llm):
        normalizer = EntityNormalizer()
        normalized_claims = await normalizer.normalize_entities(sample_claims)
        
        # Verify claim 1 entities (ASA and acetylsalicylic acid) both resolved to Aspirin / MeSH:D001241
        assert normalized_claims[0].entities[0].text == "Aspirin"
        assert normalized_claims[0].entities[0].canonical_id == "MeSH:D001241"
        assert normalized_claims[0].entities[1].text == "Aspirin"
        assert normalized_claims[0].entities[1].canonical_id == "MeSH:D001241"
        
        # Verify claim 2 entities
        # aspirin -> Aspirin / MeSH:D001241
        assert normalized_claims[1].entities[0].text == "Aspirin"
        assert normalized_claims[1].entities[0].canonical_id == "MeSH:D001241"
        
        # glucophage -> Metformin / MeSH:D008687
        assert normalized_claims[1].entities[1].text == "Metformin"
        assert normalized_claims[1].entities[1].canonical_id == "MeSH:D008687"
        
        # insulin -> Insulin / MeSH:D007328 (from LLM)
        assert normalized_claims[1].entities[2].text == "Insulin"
        assert normalized_claims[1].entities[2].canonical_id == "MeSH:D007328"
        
        # Verify LLM was called exactly once for the single unresolved entity 'insulin'
        mock_llm.generate_structured.assert_called_once()
        call_kwargs = mock_llm.generate_structured.call_args[1]
        assert "insulin" in call_kwargs["prompt"]
        # And fast-lookup entities were NOT sent to the LLM as input
        assert "glucophage" not in call_kwargs["prompt"]



@pytest.mark.asyncio
async def test_normalize_entities_llm_fallback_failure(sample_claims):
    # Test that if the LLM fallback fails, original names are kept for unresolved entities
    mock_llm = MagicMock()
    mock_llm.model_name = "mock-llm"
    mock_llm.generate_structured = AsyncMock(side_effect=Exception("API failure"))
    
    with patch("src.entity.normalizer.get_llm", return_value=mock_llm):
        normalizer = EntityNormalizer()
        normalized_claims = await normalizer.normalize_entities(sample_claims)
        
        # Local lookup ones are still resolved
        assert normalized_claims[0].entities[0].text == "Aspirin"
        assert normalized_claims[1].entities[1].text == "Metformin"
        
        # Insulin (which failed LLM resolution) kept its original text and None canonical_id
        assert normalized_claims[1].entities[2].text == "insulin"
        assert normalized_claims[1].entities[2].canonical_id is None


@pytest.mark.asyncio
async def test_normalizer_loads_custom_synonym_map(tmp_path):
    # Create a custom synonym map file
    custom_map = {
        "advil": ["Ibuprofen", "MeSH:D007052"],
        "tylenol": ["Acetaminophen", "MeSH:D000082"]
    }
    custom_file = tmp_path / "custom_synonyms.json"
    custom_file.write_text(json.dumps(custom_map))

    normalizer = EntityNormalizer(synonym_map_path=str(custom_file))
    
    # Assert custom map loaded correctly
    assert normalizer.synonym_map["advil"] == ("Ibuprofen", "MeSH:D007052")
    assert normalizer.synonym_map["tylenol"] == ("Acetaminophen", "MeSH:D000082")

    # Assert fallback default items are NOT in custom map (since it loaded custom only)
    assert "aspirin" not in normalizer.synonym_map


def test_normalizer_fallback_on_missing_file():
    # Instantiate with non-existent path
    normalizer = EntityNormalizer(synonym_map_path="non_existent_file.json")
    
    # Should fall back to default map
    assert "aspirin" in normalizer.synonym_map
    assert normalizer.synonym_map["aspirin"] == ("Aspirin", "MeSH:D001241")


def test_normalizer_fallback_on_invalid_json(tmp_path):
    # Create an invalid JSON file
    invalid_file = tmp_path / "invalid.json"
    invalid_file.write_text("invalid json content {")

    normalizer = EntityNormalizer(synonym_map_path=str(invalid_file))
    
    # Should fall back to default map
    assert "aspirin" in normalizer.synonym_map
    assert normalizer.synonym_map["aspirin"] == ("Aspirin", "MeSH:D001241")