didi-lot1-ai/ai_platform/modules/web/tests/test_vision_client.py

"""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"