"""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, {})