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