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

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

import pytest

from vmk_data_collector.core.exceptions import (
    OllamaFatalError,
    OllamaRetryableError,
)
from vmk_data_collector.domain.entities import (
    AiEnrichmentResult,
    NormalizedProperty,
)
from vmk_data_collector.services.ai_enricher import AiEnricher


@pytest.fixture
def enricher(mock_ollama_client: AsyncMock) -> AiEnricher:
    return AiEnricher(client=mock_ollama_client)


@pytest.fixture
def normalized_property() -> NormalizedProperty:
    return NormalizedProperty(
        property_type="apartment",
        deal_type="sale",
        title="Test Title",
        description="Test Description",
        price=100000,
        currency="UAH",
        total_area=50,
        rooms_count=2,
        floor=3,
        floors_total=9,
        city="Kyiv",
        address_raw="Kyiv",
    )


class TestMockMode:
    @pytest.mark.asyncio
    async def test_returns_mock_when_enabled(
        self,
        enricher: AiEnricher,
        normalized_property: NormalizedProperty,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        monkeypatch.setattr(
            "vmk_data_collector.services.ai_enricher.settings.ollama_mock",
            True,
        )
        result = await enricher.enrich(normalized_property, {})
        assert isinstance(result, AiEnrichmentResult)
        assert result.classification == "жилая_недвижимость"


class TestHappyPath:
    @pytest.mark.asyncio
    async def test_parses_json_response(
        self,
        enricher: AiEnricher,
        mock_ollama_client: AsyncMock,
        normalized_property: NormalizedProperty,
    ) -> None:
        mock_ollama_client.chat.return_value = {
            "message": {
                "content": json.dumps(
                    {
                        "extracted_features": {},
                        "price_assessment": {
                            "estimated_market_price": 120000,
                            "price_reasonableness": "на уровне рынка",
                            "currency": "UAH",
                        },
                        "listing_quality_score": 7,
                        "reliability_rating": 4,
                        "sentiment_score": 0.5,
                        "classification": "жилая_недвижимость",
                        "image_analysis_results": {},
                        "generated_description": "GD",
                        "summary": "S",
                        "model_version": "v1",
                        "processing_time_ms": 100,
                    }
                )
            }
        }
        result = await enricher.enrich(normalized_property, {})
        assert isinstance(result, AiEnrichmentResult)
        assert result.listing_quality_score == 7


class TestErrorHandling:
    @pytest.mark.asyncio
    async def test_retryable_error_propagates(
        self,
        enricher: AiEnricher,
        mock_ollama_client: AsyncMock,
        normalized_property: NormalizedProperty,
    ) -> None:
        mock_ollama_client.chat.side_effect = OllamaRetryableError(
            "transient"
        )
        with pytest.raises(OllamaRetryableError):
            await enricher.enrich(normalized_property, {})

    @pytest.mark.asyncio
    async def test_fatal_error_propagates(
        self,
        enricher: AiEnricher,
        mock_ollama_client: AsyncMock,
        normalized_property: NormalizedProperty,
    ) -> None:
        mock_ollama_client.chat.side_effect = OllamaFatalError("fatal")
        with pytest.raises(OllamaFatalError):
            await enricher.enrich(normalized_property, {})

    @pytest.mark.asyncio
    async def test_validation_error_returns_none(
        self,
        enricher: AiEnricher,
        normalized_property: NormalizedProperty,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        import pydantic

        def bad_init(**_kwargs):
            raise pydantic.ValidationError.from_exception_data(
                "AiEnrichmentResult", []
            )

        monkeypatch.setattr(
            "vmk_data_collector.services.ai_enricher.AiEnrichmentResult",
            bad_init,
        )
        result = await enricher.enrich(normalized_property, {})
        assert result is None

    @pytest.mark.asyncio
    async def test_unexpected_error_propagates(
        self,
        enricher: AiEnricher,
        mock_ollama_client: AsyncMock,
        normalized_property: NormalizedProperty,
    ) -> None:
        mock_ollama_client.chat.side_effect = RuntimeError("unexpected")
        with pytest.raises(RuntimeError):
            await enricher.enrich(normalized_property, {})