Newer
Older
vmk-360-data_collector / tests / unit / test_ai_normalizer.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy 1 day ago 4 KB fix: code review critical and high issues
"""Unit tests for AiNormalizer."""

import json
from typing import Any
from unittest.mock import AsyncMock

import pydantic
import pytest

from vmk_data_collector.core.exceptions import (
    OllamaFatalError,
    OllamaRetryableError,
)
from vmk_data_collector.schemas.ai_response import AiNormalizerResponse
from vmk_data_collector.services.ai_normalizer import AiNormalizer


@pytest.fixture
def normalizer(mock_ollama_client: AsyncMock) -> AiNormalizer:
    return AiNormalizer(client=mock_ollama_client)


class TestMockMode:
    @pytest.mark.asyncio
    async def test_returns_mock_when_enabled(
        self,
        normalizer: AiNormalizer,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        monkeypatch.setattr(
            "vmk_data_collector.services.ai_normalizer.settings.ollama_mock",
            True,
        )
        result = await normalizer.normalize({"title": "any"})
        assert result.is_real_estate is True
        assert result.normalized["title"] == "Mock Title"


class TestHappyPath:
    @pytest.mark.asyncio
    async def test_parses_json_response(
        self,
        normalizer: AiNormalizer,
        mock_ollama_client: AsyncMock,
    ) -> None:
        payload: dict[str, Any] = {
            "title": "1-комнатная квартира",
            "price": 50000,
        }
        mock_ollama_client.chat.return_value = {
            "message": {
                "content": json.dumps(
                    {
                        "is_real_estate": True,
                        "reason": None,
                        "normalized": {
                            "property_type": "apartment",
                            "deal_type": "sale",
                            "title": "1-комнатная",
                            "description": "desc",
                            "price": 50000,
                            "currency": "UAH",
                            "total_area": 30,
                            "rooms_count": 1,
                            "floor": 2,
                            "floors_total": 5,
                            "city": "Odessa",
                            "address_raw": "Odessa",
                            "images": [],
                            "custom_fields": {},
                        },
                    }
                )
            }
        }
        result = await normalizer.normalize(payload)
        assert isinstance(result, AiNormalizerResponse)
        assert result.is_real_estate is True
        assert result.normalized["city"] == "Odessa"

    @pytest.mark.asyncio
    async def test_build_text_is_json_and_ignores_images(
        self, normalizer: AiNormalizer
    ) -> None:
        text = normalizer._build_text(
            {
                "title": "T",
                "description": "D",
                "price": "100",
                "url": "http://x",
                "images": ["a.jpg"],
                "extra": "value",
            }
        )
        assert '"title": "T"' in text
        assert '"extra": "value"' in text
        assert "images" not in text


class TestErrorHandling:
    @pytest.mark.asyncio
    async def test_retryable_error_propagates(
        self,
        normalizer: AiNormalizer,
        mock_ollama_client: AsyncMock,
    ) -> None:
        mock_ollama_client.chat.side_effect = OllamaRetryableError(
            "transient"
        )
        with pytest.raises(OllamaRetryableError):
            await normalizer.normalize({})

    @pytest.mark.asyncio
    async def test_fatal_error_propagates(
        self,
        normalizer: AiNormalizer,
        mock_ollama_client: AsyncMock,
    ) -> None:
        mock_ollama_client.chat.side_effect = OllamaFatalError("fatal")
        with pytest.raises(OllamaFatalError):
            await normalizer.normalize({})

    @pytest.mark.asyncio
    async def test_validation_error_returns_invalid(
        self,
        normalizer: AiNormalizer,
        mock_ollama_client: AsyncMock,
    ) -> None:
        mock_ollama_client.chat.return_value = {
            "message": {"content": '{"is_real_estate": []}'}
        }
        result = await normalizer.normalize({})
        assert result.is_real_estate is False
        assert "validation error" in result.reason.lower()

    @pytest.mark.asyncio
    async def test_unexpected_error_propagates(
        self,
        normalizer: AiNormalizer,
        mock_ollama_client: AsyncMock,
    ) -> None:
        mock_ollama_client.chat.side_effect = RuntimeError("unexpected")
        with pytest.raises(RuntimeError):
            await normalizer.normalize({})