From 8280e02aed21b0f1e52c3fc4333fed00a11a1ade Mon Sep 17 00:00:00 2001
From: Dmitry Afanasyev <Balshbox@gmail.com>
Date: Tue, 30 Aug 2022 09:47:40 +0300
Subject: [PATCH] add busses list

---
 app/core/bot.py                | 12 ++++++++++++
 app/core/parse_web.py          | 28 ++++++++++++++--------------
 tests/bot/test_bot_selenoid.py |  4 ++++
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/app/core/bot.py b/app/core/bot.py
index 89346f2..490fa66 100644
--- a/app/core/bot.py
+++ b/app/core/bot.py
@@ -50,6 +50,10 @@ class TransportBot:
             driver=driver,
             url='https://yandex.ru/maps/213/moscow/stops/stop__9640740/?ll=37.527924%2C55.823470&tab=overview&z=21',
             message='Остановка Б. Академическая ул, д. 15',
+            buses=[
+                '300',
+                'т19',
+            ],
         )
 
         return await TransportBot.bot.send_message(
@@ -66,6 +70,10 @@ class TransportBot:
             driver=driver,
             url='https://yandex.ru/maps/213/moscow/stops/stop__9640288/?ll=37.505402%2C55.800214&tab=overview&z=21',
             message='Остановка Улица Алабяна',
+            buses=[
+                '300',
+                'т19',
+            ],
         )
 
         return await TransportBot.bot.send_message(
@@ -91,6 +99,10 @@ class TransportBot:
             driver=driver,
             url='https://yandex.ru/maps/213/moscow/stops/stop__9640740/?ll=37.527924%2C55.823470&tab=overview&z=21',
             message='Остановка Б. Академическая ул, д. 15',
+            buses=[
+                '300',
+                'т19',
+            ],
         )
         await asyncio.gather(
             *[
diff --git a/app/core/parse_web.py b/app/core/parse_web.py
index 55e2a6a..31d1033 100644
--- a/app/core/parse_web.py
+++ b/app/core/parse_web.py
@@ -57,7 +57,11 @@ class WebParser:
 
     @staticmethod
     def parse_yandex_maps(
-        url: str, message: str, driver: RemoteWebDriver | None = None
+        *,
+        url: str,
+        message: str,
+        buses: list[str],
+        driver: RemoteWebDriver | None = None,
     ) -> str:
         if not driver:
             logger.error('Driver is not configured')
@@ -81,26 +85,22 @@ class WebParser:
                         by='class name',
                         value='masstransit-prognoses-view__title-text',
                     )
-                    match bus.text:
-                        case "300":
-                            bus_arrival["Автобус 300"] = (
-                                bus_arrival_time.text if bus_arrival_time else None
-                            )
-                        case "т19":
-                            bus_arrival["Автобус Т19"] = (
-                                bus_arrival_time.text if bus_arrival_time else None
-                            )
+                    bus_arrival[bus.text] = (
+                        bus_arrival_time.text if bus_arrival_time else None
+                    )
         except NoSuchElementException:
             pass
         except StaleElementReferenceException:
             pass
 
-        if not any(bus_arrival.values()):
-            return 'Автобусов 300 или Т19 не найдено. \n\nСмотри на карте :)'
+        if not any([bus_arrival.get(bus_name) for bus_name in buses]):
+            return f'Автобусов {", ".join(buses)} не найдено. \n\nСмотри на карте :)'
 
         answer = f'{message}\n\n'
-        for bus_name, arrival_time in bus_arrival.items():
-            answer += f'{bus_name} - {arrival_time}\n'
+        for bus_name in buses:
+            arrival_time = bus_arrival.get(bus_name)
+            if arrival_time:
+                answer += f'Автобус {bus_name} - {arrival_time}\n'
         return answer
 
     @staticmethod
diff --git a/tests/bot/test_bot_selenoid.py b/tests/bot/test_bot_selenoid.py
index 4ab0cc4..815b07d 100644
--- a/tests/bot/test_bot_selenoid.py
+++ b/tests/bot/test_bot_selenoid.py
@@ -69,5 +69,9 @@ async def test_selenoid_parse_yandex() -> None:
         driver=driver,
         url='https://yandex.ru/maps/213/moscow/stops/stop__9640740/?ll=37.527924%2C55.823470&tab=overview&z=21',
         message='Остановка Б. Академическая ул, д. 15',
+        buses=[
+            '300',
+            'т19',
+        ],
     )
     assert len(text) > 0