Newer
Older
vmk-360-data_collector / tests / conftest.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy 1 day ago 3 KB fix: code review critical and high issues
"""Pytest fixtures and helpers."""

from collections.abc import AsyncGenerator
from typing import Any
from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest
from fastapi import FastAPI
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from sqlalchemy.ext.asyncio import AsyncSession

from vmk_data_collector.api.deps import get_db
from vmk_data_collector.api.v1.router_health import router as health_router
from vmk_data_collector.api.v1.router_properties import router as properties_router
from vmk_data_collector.core.exceptions import (
    AIProcessingError,
    AppError,
    NotRealEstateError,
    ValidationError,
)
from vmk_data_collector.core.limiter import limiter
from vmk_data_collector.main import (
    ai_processing_error_handler,
    app_error_handler,
    not_real_estate_handler,
    validation_error_handler,
)


@pytest.fixture
def mock_async_session() -> AsyncMock:
    """Mock SQLAlchemy AsyncSession with common methods."""
    session = AsyncMock(spec=AsyncSession)
    session.__aenter__ = AsyncMock(return_value=session)
    session.__aexit__ = AsyncMock(return_value=None)

    result_mock = MagicMock()
    session.execute.return_value = result_mock
    return session


@pytest.fixture
def mock_session_factory(mock_async_session: AsyncMock) -> MagicMock:
    """Mock async_sessionmaker that yields mock_async_session."""
    factory = MagicMock()
    factory.return_value = mock_async_session
    return factory


@pytest.fixture
def mock_ollama_client() -> AsyncMock:
    """Mock OllamaClient with a default successful chat response."""
    client = AsyncMock()
    client.chat.return_value = {
        "message": {
            "content": (
                '{"is_real_estate": true, "reason": null, '
                '"normalized": {"property_type": "apartment", '
                '"deal_type": "sale", "title": "Test", '
                '"description": "Desc", "price": 100000, '
                '"currency": "UAH", "total_area": 50, '
                '"rooms_count": 2, "floor": 3, '
                '"floors_total": 9, "city": "Kyiv", '
                '"address_raw": "Kyiv", "images": [], '
                '"custom_fields": {}}}'
            )
        }
    }
    return client


@pytest.fixture
def fastapi_app(mock_async_session: AsyncMock) -> FastAPI:
    """FastAPI app for integration tests (no lifespan, no worker)."""
    app = FastAPI()
    app.state.limiter = limiter
    app.add_exception_handler(
        RateLimitExceeded, _rate_limit_exceeded_handler
    )
    app.add_exception_handler(AppError, app_error_handler)
    app.add_exception_handler(ValidationError, validation_error_handler)
    app.add_exception_handler(NotRealEstateError, not_real_estate_handler)
    app.add_exception_handler(AIProcessingError, ai_processing_error_handler)
    app.include_router(health_router, prefix="/api/v1")
    app.include_router(properties_router, prefix="/api/v1")
    app.dependency_overrides[get_db] = lambda: mock_async_session
    return app


@pytest.fixture
async def async_client(
    fastapi_app: FastAPI,
) -> AsyncGenerator[httpx.AsyncClient, None]:
    """Async HTTP client wired to the test FastAPI app."""
    transport = httpx.ASGITransport(app=fastapi_app)
    async with httpx.AsyncClient(
        transport=transport, base_url="http://test"
    ) as client:
        yield client