tests/test_knowledge_base.py

"""
Tests for Knowledge Base module
"""

import pytest
from pathlib import Path
import tempfile
from knowledge_base import MedicalTermManager, load_terms_from_file, save_terms_to_file


class TestMedicalTermManager:
    """Tests for MedicalTermManager"""
    
    def test_initialization(self):
        """Test basic initialization"""
        manager = MedicalTermManager()
        assert manager is not None
        assert isinstance(manager.terms, set)
    
    def test_add_term(self):
        """Test adding a term"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        
        assert "МРТ" in manager.terms
    
    def test_normalize_term(self):
        """Test term normalization"""
        manager = MedicalTermManager()
        
        # Test removing extra spaces
        normalized = manager.normalize_term("  МРТ  ")
        assert normalized == "МРТ"
        
        normalized = manager.normalize_term("Т1-ВИ   режим")
        assert normalized == "Т1-ВИ режим"
    
    def test_categorization(self):
        """Test automatic categorization"""
        manager = MedicalTermManager()
        
        # Add various terms
        manager.add_term("МРТ")
        manager.add_term("Т1-ВИ")
        manager.add_term("L4-L5")
        manager.add_term("грыжа")
        manager.add_term("гиперинтенсивный")
        
        # Check categories
        assert len(manager.categories['imaging']) > 0
        assert len(manager.categories['sequences']) > 0
        assert len(manager.categories['anatomy']) > 0
        assert len(manager.categories['pathology']) > 0
        assert len(manager.categories['modifiers']) > 0
    
    def test_search_term(self):
        """Test term search"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("КТ")
        manager.add_term("МСКТ")
        
        # Search for "МРТ"
        results = manager.search_term("МРТ")
        assert "МРТ" in results
        
        # Search for partial match
        results = manager.search_term("КТ")
        assert len(results) >= 2  # Should match "КТ" and "МСКТ"
    
    def test_get_category_terms(self):
        """Test getting terms by category"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("КТ")
        
        imaging_terms = manager.get_category_terms('imaging')
        assert len(imaging_terms) >= 2
    
    def test_get_all_terms(self):
        """Test getting all terms"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("КТ")
        manager.add_term("грыжа")
        
        all_terms = manager.get_all_terms()
        assert len(all_terms) == 3
        assert "МРТ" in all_terms
    
    def test_get_terms_as_text(self):
        """Test getting terms as formatted text"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("КТ")
        
        text = manager.get_terms_as_text()
        assert "МРТ" in text
        assert "КТ" in text
        assert ", " in text  # Default separator
    
    def test_validate_transcription(self):
        """Test transcription validation"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("грыжа")
        
        transcription = "Проведено МРТ. Обнаружена грыжа диска."
        
        validation = manager.validate_transcription(transcription)
        assert validation['count'] == 2
        assert "МРТ" in validation['found_terms']
        assert "грыжа" in validation['found_terms']
    
    def test_statistics(self):
        """Test statistics generation"""
        manager = MedicalTermManager()
        manager.add_term("МРТ")
        manager.add_term("Т1-ВИ")
        manager.add_term("грыжа")
        
        stats = manager.get_statistics()
        assert stats['total_terms'] == 3
        assert 'categories' in stats


class TestTermLoader:
    """Tests for term loading/saving functions"""
    
    def test_load_terms_from_file(self):
        """Test loading terms from a file"""
        # Create temporary file
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt', encoding='utf-8') as f:
            f.write("МРТ, КТ, МСКТ, грыжа")
            temp_file = Path(f.name)
        
        try:
            terms = load_terms_from_file(temp_file)
            assert len(terms) == 4
            assert "МРТ" in terms
            assert "КТ" in terms
        finally:
            temp_file.unlink()
    
    def test_save_terms_to_file(self):
        """Test saving terms to a file"""
        terms = ["МРТ", "КТ", "МСКТ"]
        
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f:
            temp_file = Path(f.name)
        
        try:
            result = save_terms_to_file(terms, temp_file)
            assert result is True
            
            # Verify file contents
            content = temp_file.read_text(encoding='utf-8')
            assert "МРТ" in content
            assert "КТ" in content
            assert "МСКТ" in content
        finally:
            temp_file.unlink()


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