import unittest
from unittest.mock import patch, MagicMock
from pygeai.core.omni.managers import OmniParserManager
from pygeai.core.omni.models import OmniParserConfiguration
from pygeai.core.omni.responses import OmniParserResponse
from pygeai.core.common.exceptions import APIError


class TestOmniParserManager(unittest.TestCase):

    def setUp(self):
        self.manager = OmniParserManager()
        self.config = OmniParserConfiguration(
            file_path="test.pdf",
            start_page=1,
            end_page=5
        )

    @patch('pygeai.core.omni.managers.OmniParserClient')
    def test_process_file_success(self, mock_client_class):
        mock_client = MagicMock()
        mock_client_class.return_value = mock_client

        mock_response_data = {
            "status": "success",
            "parts": [{"text": "Sample", "type": "Text"}],
            "requestId": "123",
            "usage": {"total_cost": 0.01}
        }
        mock_client.process_file.return_value = mock_response_data

        manager = OmniParserManager()
        result = manager.process_file(self.config)

        self.assertIsInstance(result, OmniParserResponse)
        self.assertEqual(result.status, "success")

    @patch('pygeai.core.omni.managers.OmniParserClient')
    def test_process_file_with_error(self, mock_client_class):
        mock_client = MagicMock()
        mock_client_class.return_value = mock_client

        mock_response_data = {
            "error": {
                "id": 100,
                "code": "PROCESSING_ERROR",
                "message": "File processing failed"
            }
        }
        mock_client.process_file.return_value = mock_response_data

        manager = OmniParserManager()
        with self.assertRaises(APIError):
            manager.process_file(self.config)
