Добавлен yandex geosuggest

tests/services/test_address_suggestions.py

@@ -15,6 +15,9 @@ from app.services.aggregator import (
     UnsupportedAddressSuggestionCountryError,
 )
 
+_DADATA_COUNTRIES = ("RU", "BY", "KZ")
+_YANDEX_COUNTRIES = ("AM", "AZ", "KG", "MD", "TJ", "TM", "UZ")
+
 
 class StubAddressSuggestionProvider:
     def __init__(
@@ -64,7 +67,20 @@ def _make_suggestion(
     )
 
 
-def test_suggest_addresses_routes_ru_to_dadata() -> None:
+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=[
@@ -77,16 +93,20 @@ def test_suggest_addresses_routes_ru_to_dadata() -> None:
             )
         ],
     )
-    europe = StubAddressSuggestionProvider(
-        "europe",
+    yandex_geosuggest = StubAddressSuggestionProvider(
+        "yandex_geosuggest",
         response=[_make_suggestion("Yerevan, Tumanyan 1")],
     )
     service = AggregatorService(
         providers=[],
-        address_suggestion_providers=[dadata, europe],
-        address_suggestion_country_to_provider={"RU": "dadata", "AM": "europe"},
+        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",
     )
-    request = _make_request(country_code="RU", city="Moscow", query="Khabarovskaya")
 
     result = asyncio.run(service.suggest_addresses(request))
 
@@ -100,30 +120,33 @@ def test_suggest_addresses_routes_ru_to_dadata() -> None:
         )
     ]
     assert dadata.calls == [request]
-    assert europe.calls == []
+    assert yandex_geosuggest.calls == []
 
 
-def test_suggest_addresses_routes_second_provider_by_country_mapping() -> None:
+@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")],
     )
-    europe = StubAddressSuggestionProvider(
-        "europe",
+    yandex_geosuggest = StubAddressSuggestionProvider(
+        "yandex_geosuggest",
         response=[_make_suggestion("Yerevan, Tumanyan 1")],
     )
     service = AggregatorService(
         providers=[],
-        address_suggestion_providers=[dadata, europe],
-        address_suggestion_country_to_provider={"RU": "dadata", "AM": "europe"},
+        address_suggestion_providers=[dadata, yandex_geosuggest],
+        address_suggestion_country_to_provider=_make_country_mapping(),
     )
-    request = _make_request(country_code="AM", city="Yerevan", query="Tumanyan")
+    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 europe.calls == [request]
+    assert yandex_geosuggest.calls == [request]
 
 
 def test_suggest_addresses_raises_for_unsupported_country() -> None:
@@ -155,7 +178,7 @@ def test_suggest_addresses_raises_for_unregistered_provider_mapping() -> None:
                 response=[_make_suggestion("Moscow, Lenina 1")],
             )
         ],
-        address_suggestion_country_to_provider={"AM": "europe"},
+        address_suggestion_country_to_provider={"AM": "yandex_geosuggest"},
     )
 
     with pytest.raises(UnsupportedAddressSuggestionCountryError):