tests/test_ocr_handler.py - master - alsoalgo/CTF

tests/test_ocr_handler.py
208 строк8 KB
alsoalgo
some improvements
14 дек 2025, 22:23
14 дек 2025, 22:2337d2ec0
"""
Тесты для OCRHandler
"""
import pytest
import os
import tempfile
from unittest.mock import Mock, patch, MagicMock
import json
import pandas as pd

from src.components.ocr_handler import OCRHandler


class TestOCRHandler:
    """Тесты для OCRHandler"""
    
    def setup_method(self):
        """Настройка перед каждым тестом"""
        # Используем тестовый API ключ
        self.test_api_key = "test_api_key_12345"
        self.handler = OCRHandler(api_key=self.test_api_key)
    
    def test_initialization_with_api_key(self):
        """Тест инициализации с API ключом"""
        handler = OCRHandler(api_key="test_key")
        assert handler.api_key == "test_key"
        assert handler.base_url == "https://openrouter.ai/api/v1/chat/completions"
    
    def test_initialization_without_api_key(self):
        """Тест инициализации без API ключа"""
        with patch.dict(os.environ, {}, clear=True):
            handler = OCRHandler()
            assert handler.api_key is None
    
    def test_initialization_from_env(self):
        """Тест инициализации из переменной окружения"""
        with patch.dict(os.environ, {'OPENROUTER_API_KEY': 'env_key'}):
            handler = OCRHandler()
            assert handler.api_key == 'env_key'
    
    def test_encode_file_to_base64_success(self):
        """Тест успешного кодирования файла в base64"""
        # Создаем временный файл
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as tmp:
            tmp.write("test content")
            tmp_path = tmp.name
        
        try:
            result = self.handler._encode_file_to_base64(tmp_path)
            assert result is not None
            assert result.startswith('data:')
            assert 'base64,' in result
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_encode_file_to_base64_nonexistent(self):
        """Тест кодирования несуществующего файла"""
        result = self.handler._encode_file_to_base64("/nonexistent/file.txt")
        assert result is None
    
    def test_encode_file_to_base64_pdf_mime_type(self):
        """Тест правильного MIME типа для PDF"""
        # Создаем временный PDF файл (просто с расширением .pdf)
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf content")
            tmp_path = tmp.name
        
        try:
            result = self.handler._encode_file_to_base64(tmp_path)
            assert result is not None
            assert 'application/pdf' in result
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    @patch('src.components.ocr_handler.requests.post')
    def test_extract_text_with_ocr_success(self, mock_post):
        """Тест успешного извлечения текста через OCR"""
        # Мокаем ответ API
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            'choices': [{
                'message': {
                    'content': 'Extracted text from document'
                }
            }]
        }
        mock_post.return_value = mock_response
        
        # Создаем временный файл
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf")
            tmp_path = tmp.name
        
        try:
            result = self.handler.extract_text_with_ocr(tmp_path)
            assert result == 'Extracted text from document'
            assert mock_post.called
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    @patch('src.components.ocr_handler.requests.post')
    def test_extract_text_with_ocr_api_error(self, mock_post):
        """Тест обработки ошибки API"""
        # Мокаем ошибку API
        mock_response = Mock()
        mock_response.status_code = 400
        mock_response.json.return_value = {'error': 'Bad request'}
        mock_post.return_value = mock_response
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf")
            tmp_path = tmp.name
        
        try:
            result = self.handler.extract_text_with_ocr(tmp_path)
            assert result is None
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_extract_text_with_ocr_no_api_key(self):
        """Тест извлечения текста без API ключа"""
        handler = OCRHandler(api_key=None)
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf")
            tmp_path = tmp.name
        
        try:
            result = handler.extract_text_with_ocr(tmp_path)
            assert result is None
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_extract_text_with_ocr_nonexistent_file(self):
        """Тест извлечения текста из несуществующего файла"""
        result = self.handler.extract_text_with_ocr("/nonexistent/file.pdf")
        assert result is None
    
    @patch('src.components.ocr_handler.requests.post')
    def test_extract_structured_data_success(self, mock_post):
        """Тест успешного извлечения структурированных данных"""
        # Мокаем ответ API с таблицей
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            'choices': [{
                'message': {
                    'content': '| Field | Value |\n|--------|-------|\n| Supplier | Test Supplier |\n| Amount | 1000 |'
                }
            }]
        }
        mock_post.return_value = mock_response
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf")
            tmp_path = tmp.name
        
        try:
            result = self.handler.extract_structured_data(tmp_path)
            # Может быть None или DataFrame
            assert result is None or isinstance(result, pd.DataFrame)
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    @patch('src.components.ocr_handler.requests.post')
    def test_extract_text_with_ocr_and_annotations(self, mock_post):
        """Тест извлечения текста с аннотациями"""
        # Мокаем ответ API с аннотациями
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            'choices': [{
                'message': {
                    'content': 'Extracted text',
                    'annotations': {
                        'text_blocks': [
                            {
                                'text': 'Test',
                                'bbox': [10, 20, 100, 30]
                            }
                        ]
                    }
                }
            }]
        }
        mock_post.return_value = mock_response
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf")
            tmp_path = tmp.name
        
        try:
            # Используем extract_text_with_ocr с return_annotations=True
            result = self.handler.extract_text_with_ocr(tmp_path, return_annotations=True)
            # Может вернуть (text, annotations) или None
            if result is not None:
                assert isinstance(result, tuple)
                assert len(result) == 2
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)