g2s-aggregator/tests/services/test_address_suggestions.py

import asyncio

import pytest

from app.adapters.address_suggestions.base import (
    AddressSuggestionClientError,
    AddressSuggestionRequestError,
)
from app.schemas.request import AddressSuggestRequest
from app.schemas.response import AddressSuggestion
from app.services.aggregator import (
    AddressSuggestionsUnavailableError,
    AggregatorService,
    InvalidAddressSuggestRequestError,
    UnsupportedAddressSuggestionCountryError,
)

_DADATA_COUNTRIES = ("RU", "BY", "KZ")
_YANDEX_COUNTRIES = ("AM", "AZ", "KG", "MD", "TJ", "TM", "UZ")


class StubAddressSuggestionProvider:
    def __init__(
        self,
        name: str,
        *,
        response: list[AddressSuggestion] | None = None,
        error: Exception | None = None,
    ) -> None:
        self.name = name
        self._response = response if response is not None else []
        self._error = error
        self.calls: list[AddressSuggestRequest] = []

    async def suggest(self, request: AddressSuggestRequest) -> list[AddressSuggestion]:
        self.calls.append(request)
        if self._error is not None:
            raise self._error
        return self._response


def _make_request(**overrides: object) -> AddressSuggestRequest:
    payload: dict[str, object] = {
        "country_code": "RU",
        "city": "Moscow",
        "query": "Lenina",
        "limit": 5,
    }
    payload.update(overrides)
    return AddressSuggestRequest(**payload)


def _make_suggestion(
    address: str,
    *,
    street: str | None = None,
    house: str | None = None,
    flat: str | None = None,
    postal_code: str | None = None,
) -> AddressSuggestion:
    return AddressSuggestion(
        address=address,
        street=street,
        house=house,
        flat=flat,
        postal_code=postal_code,
    )


def _make_country_mapping() -> dict[str, str]:
    return {
        **{country_code: "dadata" for country_code in _DADATA_COUNTRIES},
        **{
            country_code: "yandex_geosuggest"
            for country_code in _YANDEX_COUNTRIES
        },
    }


@pytest.mark.parametrize("country_code", _DADATA_COUNTRIES)
def test_suggest_addresses_routes_dadata_countries_to_dadata(
    country_code: str,
) -> None:
    dadata = StubAddressSuggestionProvider(
        "dadata",
        response=[
            _make_suggestion(
                "107241, Moscow, Khabarovskaya 1",
                street="Khabarovskaya",
                house="1",
                flat="25",
                postal_code="107241",
            )
        ],
    )
    yandex_geosuggest = StubAddressSuggestionProvider(
        "yandex_geosuggest",
        response=[_make_suggestion("Yerevan, Tumanyan 1")],
    )
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[dadata, yandex_geosuggest],
        address_suggestion_country_to_provider=_make_country_mapping(),
    )
    request = _make_request(
        country_code=country_code,
        city="Moscow",
        query="Khabarovskaya",
    )

    result = asyncio.run(service.suggest_addresses(request))

    assert result == [
        AddressSuggestion(
            address="107241, Moscow, Khabarovskaya 1",
            street="Khabarovskaya",
            house="1",
            flat="25",
            postal_code="107241",
        )
    ]
    assert dadata.calls == [request]
    assert yandex_geosuggest.calls == []


@pytest.mark.parametrize("country_code", _YANDEX_COUNTRIES)
def test_suggest_addresses_routes_cis_countries_to_yandex_geosuggest(
    country_code: str,
) -> None:
    dadata = StubAddressSuggestionProvider(
        "dadata",
        response=[_make_suggestion("Moscow, Lenina 1")],
    )
    yandex_geosuggest = StubAddressSuggestionProvider(
        "yandex_geosuggest",
        response=[_make_suggestion("Yerevan, Tumanyan 1")],
    )
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[dadata, yandex_geosuggest],
        address_suggestion_country_to_provider=_make_country_mapping(),
    )
    request = _make_request(country_code=country_code, city="Yerevan", query="Tumanyan")

    result = asyncio.run(service.suggest_addresses(request))

    assert result == [AddressSuggestion(address="Yerevan, Tumanyan 1")]
    assert dadata.calls == []
    assert yandex_geosuggest.calls == [request]


def test_suggest_addresses_raises_for_unsupported_country() -> None:
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[
            StubAddressSuggestionProvider(
                "dadata",
                response=[_make_suggestion("Moscow, Lenina 1")],
            )
        ],
        address_suggestion_country_to_provider={"RU": "dadata"},
    )

    with pytest.raises(UnsupportedAddressSuggestionCountryError):
        asyncio.run(
            service.suggest_addresses(
                _make_request(country_code="DE", city="Berlin", query="Alexanderplatz")
            )
        )


def test_suggest_addresses_raises_for_unregistered_provider_mapping() -> None:
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[
            StubAddressSuggestionProvider(
                "dadata",
                response=[_make_suggestion("Moscow, Lenina 1")],
            )
        ],
        address_suggestion_country_to_provider={"AM": "yandex_geosuggest"},
    )

    with pytest.raises(UnsupportedAddressSuggestionCountryError):
        asyncio.run(
            service.suggest_addresses(
                _make_request(country_code="AM", city="Yerevan", query="Tumanyan")
            )
        )


def test_suggest_addresses_maps_provider_request_error_to_invalid_request() -> None:
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[
            StubAddressSuggestionProvider(
                "dadata",
                error=AddressSuggestionRequestError("provider rejected request"),
            )
        ],
        address_suggestion_country_to_provider={"RU": "dadata"},
    )

    with pytest.raises(InvalidAddressSuggestRequestError):
        asyncio.run(service.suggest_addresses(_make_request()))


def test_suggest_addresses_maps_provider_client_error_to_unavailable() -> None:
    service = AggregatorService(
        providers=[],
        address_suggestion_providers=[
            StubAddressSuggestionProvider(
                "dadata",
                error=AddressSuggestionClientError("provider unavailable"),
            )
        ],
        address_suggestion_country_to_provider={"RU": "dadata"},
    )

    with pytest.raises(AddressSuggestionsUnavailableError):
        asyncio.run(service.suggest_addresses(_make_request()))