Livrare LOT 1 - Didi

ai_platform/modules/web/tests/test_vision_client.py

@@ -0,0 +1,210 @@
+"""Unit tests for VisionClient with mocked Playwright + httpx."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import pytest
+
+from web.config import WebSettings
+
+
+def _make_httpx_response(status_code: int, json: dict) -> httpx.Response:
+    """Create an httpx.Response with a request set (needed for raise_for_status)."""
+    response = httpx.Response(status_code, json=json)
+    response._request = httpx.Request("POST", "http://test")
+    return response
+
+
+@pytest.fixture
+def settings() -> WebSettings:
+    """Create test settings."""
+    return WebSettings(
+        searxng_base_url="http://localhost:55100",
+        llm_base_url="http://localhost:14011",
+        external_url="http://localhost:51100",
+        vision_model="test-vision",
+        vision_max_tokens=1000,
+    )
+
+
+class TestVisionClientImport:
+    """Tests for VisionClient import handling."""
+
+    def test_raises_import_error_without_playwright(
+        self, settings: WebSettings
+    ) -> None:
+        """Test that VisionClient raises ImportError without playwright."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", False):
+            from web.vision.client import VisionClient
+
+            with pytest.raises(ImportError, match="playwright"):
+                VisionClient(settings)
+
+
+class TestVisionClientExtract:
+    """Tests for VisionClient extraction with mocked browser and LLM."""
+
+    @pytest.fixture
+    def mock_page(self) -> MagicMock:
+        """Create a mock Playwright page."""
+        page = AsyncMock()
+        page.goto = AsyncMock()
+        page.title = AsyncMock(return_value="Vision Test Page")
+        page.screenshot = AsyncMock(return_value=b"\x89PNG\r\n\x00\x00")
+        page.evaluate = AsyncMock(return_value=[])
+        return page
+
+    @pytest.fixture
+    def mock_context(self, mock_page: MagicMock) -> MagicMock:
+        """Create a mock browser context."""
+        context = AsyncMock()
+        context.new_page = AsyncMock(return_value=mock_page)
+        context.close = AsyncMock()
+        return context
+
+    @pytest.fixture
+    def mock_browser(self, mock_context: MagicMock) -> MagicMock:
+        """Create a mock browser."""
+        browser = MagicMock()
+        browser.is_connected = MagicMock(return_value=True)
+        browser.new_context = AsyncMock(return_value=mock_context)
+        browser.close = AsyncMock()
+        return browser
+
+    @pytest.mark.asyncio
+    async def test_extract_single_page(
+        self, settings: WebSettings, mock_browser: MagicMock
+    ) -> None:
+        """Test extracting content from a single URL."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", True):
+            from web.schemas.vision import VisionExtractRequest
+            from web.vision.client import VisionClient
+
+            client = VisionClient(settings, browser=mock_browser)
+
+            # Mock the LLM call
+            llm_response = _make_httpx_response(
+                200,
+                json={
+                    "choices": [
+                        {
+                            "message": {
+                                "content": "Extracted text from the page screenshot."
+                            }
+                        }
+                    ],
+                    "usage": {"total_tokens": 500},
+                },
+            )
+
+            with patch.object(
+                httpx.AsyncClient, "post", new_callable=AsyncMock
+            ) as mock_post:
+                mock_post.return_value = llm_response
+
+                request = VisionExtractRequest(urls=["https://example.com/vision"])
+                response = await client.extract(request)
+
+            assert response.total_processed == 1
+            assert response.pages[0].title == "Vision Test Page"
+            assert "Extracted text" in response.pages[0].extracted_text
+
+    @pytest.mark.asyncio
+    async def test_extract_to_page_content(
+        self, settings: WebSettings, mock_browser: MagicMock
+    ) -> None:
+        """Test the convenience method extract_to_page_content."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", True):
+            from web.vision.client import VisionClient
+
+            client = VisionClient(settings, browser=mock_browser)
+
+            llm_response = _make_httpx_response(
+                200,
+                json={
+                    "choices": [{"message": {"content": "Page content via vision."}}],
+                    "usage": {"total_tokens": 300},
+                },
+            )
+
+            with patch.object(
+                httpx.AsyncClient, "post", new_callable=AsyncMock
+            ) as mock_post:
+                mock_post.return_value = llm_response
+
+                page_content = await client.extract_to_page_content(
+                    "https://example.com/test"
+                )
+
+            assert page_content.extraction_method == "vision"
+            assert page_content.text == "Page content via vision."
+            assert page_content.fallback_chain == ["http", "browse"]
+
+    @pytest.mark.asyncio
+    async def test_extract_handles_llm_failure(
+        self, settings: WebSettings, mock_browser: MagicMock
+    ) -> None:
+        """Test that LLM failure is reported properly."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", True):
+            from web.schemas.vision import VisionExtractRequest
+            from web.vision.client import VisionClient
+
+            client = VisionClient(settings, browser=mock_browser)
+
+            with patch.object(
+                httpx.AsyncClient, "post", new_callable=AsyncMock
+            ) as mock_post:
+                mock_post.side_effect = httpx.ConnectError("refused")
+
+                request = VisionExtractRequest(urls=["https://example.com/fail"])
+                response = await client.extract(request)
+
+            assert response.total_failed == 1
+            assert response.total_processed == 0
+
+    @pytest.mark.asyncio
+    async def test_close_only_closes_owned_browser(
+        self, settings: WebSettings, mock_browser: MagicMock
+    ) -> None:
+        """Test that close() only closes browser if client owns it."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", True):
+            from web.vision.client import VisionClient
+
+            client = VisionClient(settings, browser=mock_browser)
+            await client.close()
+
+            mock_browser.close.assert_not_called()
+
+
+class TestVisionClientProviderFallback:
+    """Tests for LLM provider fallback in VisionClient."""
+
+    @pytest.mark.asyncio
+    async def test_auto_provider_tries_local_first(self, settings: WebSettings) -> None:
+        """Test that auto provider tries local first."""
+        with patch("web.vision.client.HAS_PLAYWRIGHT", True):
+            from web.vision.client import VisionClient
+
+            mock_browser = MagicMock()
+            mock_browser.is_connected = MagicMock(return_value=True)
+            client = VisionClient(settings, browser=mock_browser)
+
+            llm_response = _make_httpx_response(
+                200,
+                json={
+                    "choices": [{"message": {"content": "Text"}}],
+                    "usage": {"total_tokens": 100},
+                },
+            )
+
+            with patch.object(
+                httpx.AsyncClient, "post", new_callable=AsyncMock
+            ) as mock_post:
+                mock_post.return_value = llm_response
+
+                text, _tokens, provider, _model = await client._call_vision_llm(
+                    ["base64data"], None, "auto", None
+                )
+
+            assert provider == "local"
+            assert text == "Text"