From 059880ebdc6fa947b14128055bf9442eb6af5023 Mon Sep 17 00:00:00 2001
From: Ivan Belokobylskiy <belokobylskij@gmail.com>
Date: Sat, 23 Oct 2021 12:14:32 +0300
Subject: [PATCH] Fix yandex captcha detecting (#56132)

Yandex recently switched to the new captcha page and the new version of aiomaps supports it.
The check for captcha is moved to platform setup.

Fixes #56035
---
 .../components/yandex_transport/manifest.json |  2 +-
 .../components/yandex_transport/sensor.py     | 23 +++++++++++++++----
 requirements_all.txt                          |  2 +-
 requirements_test_all.txt                     |  2 +-
 .../test_yandex_transport_sensor.py           |  1 +
 5 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/yandex_transport/manifest.json b/homeassistant/components/yandex_transport/manifest.json
index 79818f8e63e..680336fe47b 100644
--- a/homeassistant/components/yandex_transport/manifest.json
+++ b/homeassistant/components/yandex_transport/manifest.json
@@ -2,7 +2,7 @@
   "domain": "yandex_transport",
   "name": "Yandex Transport",
   "documentation": "https://www.home-assistant.io/integrations/yandex_transport",
-  "requirements": ["aioymaps==1.1.0"],
+  "requirements": ["aioymaps==1.2.1"],
   "codeowners": ["@rishatik92", "@devbis"],
   "iot_class": "cloud_polling"
 }
diff --git a/homeassistant/components/yandex_transport/sensor.py b/homeassistant/components/yandex_transport/sensor.py
index b4f7f986626..0acca753454 100644
--- a/homeassistant/components/yandex_transport/sensor.py
+++ b/homeassistant/components/yandex_transport/sensor.py
@@ -3,7 +3,7 @@
 from datetime import timedelta
 import logging
 
-from aioymaps import YandexMapsRequester
+from aioymaps import CaptchaError, YandexMapsRequester
 import voluptuous as vol
 
 from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
@@ -42,8 +42,16 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     routes = config[CONF_ROUTE]
 
     client_session = async_create_clientsession(hass, requote_redirect_url=False)
-    data = YandexMapsRequester(user_agent=USER_AGENT, client_session=client_session)
-    async_add_entities([DiscoverYandexTransport(data, stop_id, routes, name)], True)
+    ymaps = YandexMapsRequester(user_agent=USER_AGENT, client_session=client_session)
+    try:
+        await ymaps.set_new_session()
+    except CaptchaError as ex:
+        _LOGGER.error(
+            "%s. You may need to disable the integration for some time",
+            ex,
+        )
+        return
+    async_add_entities([DiscoverYandexTransport(ymaps, stop_id, routes, name)], True)
 
 
 class DiscoverYandexTransport(SensorEntity):
@@ -63,7 +71,14 @@ class DiscoverYandexTransport(SensorEntity):
         """Get the latest data from maps.yandex.ru and update the states."""
         attrs = {}
         closer_time = None
-        yandex_reply = await self.requester.get_stop_info(self._stop_id)
+        try:
+            yandex_reply = await self.requester.get_stop_info(self._stop_id)
+        except CaptchaError as ex:
+            _LOGGER.error(
+                "%s. You may need to disable the integration for some time",
+                ex,
+            )
+            return
         try:
             data = yandex_reply["data"]
         except KeyError as key_error:
diff --git a/requirements_all.txt b/requirements_all.txt
index aeff96dfa08..90e7bbd3208 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -264,7 +264,7 @@ aiovlc==0.1.0
 aiowatttime==0.1.1
 
 # homeassistant.components.yandex_transport
-aioymaps==1.1.0
+aioymaps==1.2.1
 
 # homeassistant.components.airly
 airly==1.1.0
diff --git a/requirements_test_all.txt b/requirements_test_all.txt
index 7920f130cb4..6ea83e4a02c 100644
--- a/requirements_test_all.txt
+++ b/requirements_test_all.txt
@@ -191,7 +191,7 @@ aiovlc==0.1.0
 aiowatttime==0.1.1
 
 # homeassistant.components.yandex_transport
-aioymaps==1.1.0
+aioymaps==1.2.1
 
 # homeassistant.components.airly
 airly==1.1.0
diff --git a/tests/components/yandex_transport/test_yandex_transport_sensor.py b/tests/components/yandex_transport/test_yandex_transport_sensor.py
index a727ea6e6cd..f18bd34e671 100644
--- a/tests/components/yandex_transport/test_yandex_transport_sensor.py
+++ b/tests/components/yandex_transport/test_yandex_transport_sensor.py
@@ -20,6 +20,7 @@ def mock_requester():
     """Create a mock for YandexMapsRequester."""
     with patch("aioymaps.YandexMapsRequester") as requester:
         instance = requester.return_value
+        instance.set_new_session = AsyncMock()
         instance.get_stop_info = AsyncMock(return_value=REPLY)
         yield instance