Merge pull request #2 from aitormagan/develop

Develop

serverless.yml

@@ -21,13 +21,18 @@ custom:
     definitions:
       functionErrors:
         treatMissingData: notBreaching
-
+  min_date_param:
+    dev: "/autocita/madrid/dev/min-date-info"
+    pro: "/autocita/madrid/pro/min-date-info"
   table:
     dev: "vaccine-notifications-dev"
     pro: "vaccine-notifications-pro"
   throughput:
     dev: 1
     pro: 5
+  update_centres_time:
+    dev: 300
+    pro: 1200
 
 
 provider:
@@ -46,11 +51,15 @@ provider:
         - dynamodb:PutItem
         - dynamodb:UpdateItem
         - dynamodb:DeleteItem
+        - ssm:PutParameter
+        - ssm:GetParameter
       Resource:
         - "*"
   environment:
     NOTIFICATIONS_TABLE: ${self:custom.table.${opt:stage, self:provider.stage}}
+    PARAM_MIN_DATE: ${self:custom.min_date_param.${opt:stage, self:provider.stage}}
     BOT_TOKEN: ${env:BOT_TOKEN}
+    UPDATE_CENTRES_TIME: ${self:custom.update_centres_time.${opt:stage, self:provider.stage}}
 
 functions:
   message_handler:

src/checker.py

@@ -10,7 +10,8 @@
 def notify(min_years, user_info):
     message = f"‼️ ¡Buenas noticias {user_info['name']}! El sistema de vacunación de la Comunidad de Madrid ya " \
               f"permite pedir cita a gente con {min_years} años o más. ¡🏃 Corre y pide tu cita en 🔗 " \
-              f"https://autocitavacuna.sanidadmadrid.org/ohcitacovid/!"
+              f"https://autocitavacuna.sanidadmadrid.org/ohcitacovid/!\n\n⁉️ ¿Sabes que ahora puedo darte las " \
+              f"primeras citas disponibles? Simplemente di /mindate."
     send_text(user_info["user_id"], message)
 
 

src/db.py

@@ -1,14 +1,22 @@
 import boto3
 import os
+import json
+from datetime import datetime
 
 TABLE_NAME = os.environ.get("NOTIFICATIONS_TABLE")
+MIN_DATE_PARAMETER = os.environ.get("PARAM_MIN_DATE")
 CLIENT = boto3.client('dynamodb')
+CLIENT_SSM = boto3.client('ssm')
 
 __USER_ID = "user_id"
 __NAME = "name"
 __AGE = "age"
 __NOTIFIED = "notified"
 
+__CENTRES_BY_DATE = "centres_by_date"
+__UPDATED_AT = "updated_at"
+__DATE_FORMAT = "%Y%m%d"
+
 
 def save_notification(user_id, name, age, notified=False):
     CLIENT.put_item(TableName=TABLE_NAME, Item={__USER_ID: {"S": str(user_id)}, __NAME: {"S": name},
@@ -42,4 +50,23 @@ def __parse_item(item):
         "age": int(item[__AGE]["N"]),
         "user_id": item[__USER_ID]["S"],
         "notified": item[__NOTIFIED]["BOOL"]
-    }
+    }
+
+
+def save_min_date_info(centres_by_date, update_time):
+    centres_by_date = {k.strftime(__DATE_FORMAT): v for k, v in centres_by_date.items()}
+    CLIENT_SSM.put_parameter(Name=MIN_DATE_PARAMETER, Value=json.dumps({
+        "updated_at": int(update_time.timestamp()),
+        "centres_by_date": centres_by_date
+    }), Overwrite=True, Type="String")
+
+
+def get_min_date_info():
+    try:
+        param_info = CLIENT_SSM.get_parameter(Name=MIN_DATE_PARAMETER)
+        decoded_content = json.loads(param_info["Parameter"]["Value"])
+        centres_by_date = decoded_content[__CENTRES_BY_DATE]
+        centres_by_date = {datetime.strptime(k, __DATE_FORMAT): v for k, v in centres_by_date.items()}
+        return centres_by_date, datetime.fromtimestamp(decoded_content[__UPDATED_AT])
+    except CLIENT_SSM.exceptions.ParameterNotFound:
+        return {}, None

src/message_handler.py

@@ -1,11 +1,15 @@
 import re
-from datetime import datetime
+import os
+from datetime import datetime, timedelta
+from collections import defaultdict
 from aws_lambda_powertools import Logger
+import requests
 from src import telegram_helpers
 from src import db
 from src.checker import get_min_years
 
 
+UPDATE_CENTRES_TIME = int(os.environ.get("UPDATE_CENTRES_TIME", 300))
 logger = Logger(service="vacunacovidmadridbot")
 
 
@@ -26,6 +30,9 @@ def handle_update(update):
                 answer = handle_current_age(update)
             elif message == "/subscribe":
                 answer = handle_subscribe(update)
+            elif message == "/mindate":
+                telegram_helpers.send_text(user_id, "⌛ Esto me puede llevar unos segunditos...")
+                answer = handle_min_date(update)
             else:
                 answer = handle_generic_message(update)
         except Exception:
@@ -52,7 +59,8 @@ def handle_start(update):
            f"tienes o tu año de nacimiento!\n\nOtros comandos útiles:\n-/subscribe: 🔔 Crea una suscripción para " \
            f"cuando puedas pedir cita para vacunarte\n- /help: 🙋 Muestra esta ayuda\n- /status: " \
            f"ℹ️ Muestra si ya estás suscrito\n- /cancel: 🔕 Cancela la notificación registrada\n - /currentage: " \
-           f"📆 Muestra la edad mínima con la que puedes pedir cita"
+           f"📆 Muestra la edad mínima con la que puedes pedir cita\n - /mindate: 📆 Lista las primeras citas " \
+           f"disponibles en los distintos centros de vacunación."
 
 
 def handle_cancel(update):
@@ -159,3 +167,73 @@ def get_age(user_input):
         age = datetime.now().year - age
 
     return age if age is not None and 0 <= age <= 120 else None
+
+
+def handle_min_date(_):
+
+    centres_by_date, last_update = db.get_min_date_info()
+
+    if last_update is None or (datetime.now() - last_update).seconds >= UPDATE_CENTRES_TIME:
+        centres_by_date, last_update = update_centres()
+
+    if centres_by_date:
+        message = "¡Estupendo 😊! Aquí tienes las primeras fechas disponibles en el sistema de autocita:\n\n"
+        for date in sorted(centres_by_date.keys()):
+            date_str = date.strftime("%d/%m/%Y")
+            centres = "\n".join(map(lambda x: f"- {x}", centres_by_date[date]))
+            message += f"*{date_str}*:\n{centres}\n\n"
+
+        updated_ago = int((datetime.now() - last_update).seconds / 60)
+        updated_at_msg = f"Actualizado hace {updated_ago} minutos"
+        updated_at_msg = updated_at_msg[:-1] if updated_ago == 1 else updated_at_msg
+        message += updated_at_msg
+    else:
+        message = "No he sido capaz de encontrar citas disponibles. Pruébalo de nuevo más tarde."
+
+    return message
+
+
+def update_centres():
+    centres_by_date = defaultdict(lambda: list())
+    centres = requests.post("https://autocitavacuna.sanidadmadrid.org/ohcitacovid/autocita/obtenerCentros",
+                            json={"edad_paciente": 45}, verify=False).json()
+
+    for centre in centres:
+        data_curr_month = requests.post(
+            "https://autocitavacuna.sanidadmadrid.org/ohcitacovid/autocita/obtenerHuecosMes",
+            json=get_spots_body(centre["idCentro"], centre["idPrestacion"], centre["agendas"]),
+            verify=False).json()
+        data_next_month = requests.post(
+            "https://autocitavacuna.sanidadmadrid.org/ohcitacovid/autocita/obtenerHuecosMes",
+            json=get_spots_body(centre["idCentro"], centre["idPrestacion"], centre["agendas"],
+                                month_modifier=1),
+            verify=False).json()
+
+        data = []
+        data.extend(data_curr_month if type(data_curr_month) == list else [])
+        data.extend(data_next_month if type(data_next_month) == list else [])
+        dates = [x.get("fecha") for x in data]
+        dates = [datetime.strptime(x, "%d-%m-%Y") for x in dates]
+        if dates:
+            centres_by_date[min(dates)].append(centre['descripcion'])
+
+    last_update = datetime.now()
+    db.save_min_date_info(centres_by_date, last_update)
+
+    return centres_by_date, last_update
+
+
+def get_spots_body(id_centre, id_prestacion, agendas, month_modifier=0):
+    today = datetime.now()
+    check_date = datetime(year=today.year, month=today.month, day=1) + timedelta(days=31 * month_modifier)
+
+    return {
+        "idPaciente": "1",
+        "idPrestacion": id_prestacion,
+        "agendas": agendas,
+        "idCentro": id_centre,
+        "mes": check_date.month,
+        "anyo": check_date.year,
+        "horaInicio": "08:00",
+        "horaFin": "22:00"
+    }

src/telegram_helpers.py

@@ -10,3 +10,5 @@ def send_text(chat_id, message):
     if response.status_code not in [200, 400, 403]:
         response.raise_for_status()
     return response.json()
+
+

tests/unit/test_db.py

@@ -1,10 +1,16 @@
 from unittest.mock import patch, MagicMock
+from datetime import datetime
+import json
 from src import db
 
 
 TABLE_NAME = "test-table"
 
 
+class ParameterNotFound(BaseException):
+    pass
+
+
 @patch("src.db.CLIENT")
 @patch("src.db.TABLE_NAME", TABLE_NAME)
 def test_given_user_info_when_save_notification_then_put_item_called(client_mock):
@@ -94,4 +100,59 @@ def test_when_get_non_notified_then_only_non_notified_people_returned(client_moc
     client_mock.get_paginator.return_value.paginate.return_value.build_full_result.assert_called_once_with()
 
 
+@patch("src.db.CLIENT_SSM")
+@patch("src.db.MIN_DATE_PARAMETER", "param_name")
+def test_given_centres_and_last_update_when_save_min_date_info_then_info_stored_in_ssm(ssm_mock):
+    centres_by_date = {
+        datetime(2021, 3, 6): ["hosp1", "hosp2"],
+        datetime(2021, 5, 9): ["hosp3", "hosp4"]
+    }
+    last_update = datetime.now()
+
+    db.save_min_date_info(centres_by_date, last_update)
+
+    ssm_mock.put_parameter.assert_called_once_with(Name=db.MIN_DATE_PARAMETER, Value=json.dumps({
+        "updated_at": int(last_update.timestamp()),
+        "centres_by_date": {
+            "20210306": ["hosp1", "hosp2"],
+            "20210509": ["hosp3", "hosp4"]
+        }
+    }), Overwrite=True, Type="String")
+
+
+@patch("src.db.CLIENT_SSM")
+@patch("src.db.MIN_DATE_PARAMETER", "param_name")
+def test_given_no_parameter_when_save_min_date_info_then_empty_dict_and_none_returned(ssm_mock):
+    ssm_mock.exceptions.ParameterNotFound = ParameterNotFound
+    ssm_mock.get_parameter.side_effect = ParameterNotFound("error")
+
+    centres_by_date, last_update = db.get_min_date_info()
+
+    assert centres_by_date == {}
+    assert last_update is None
+
+
+@patch("src.db.CLIENT_SSM")
+@patch("src.db.MIN_DATE_PARAMETER", "param_name")
+def test_given_parameter_when_save_min_date_info_then_info_stored_in_ssm(ssm_mock):
+    ssm_mock.exceptions.ParameterNotFound = ParameterNotFound
+    last_update = datetime.now()
+    ssm_mock.get_parameter.return_value = {
+        "Parameter": {
+            "Value": json.dumps({
+                "updated_at": int(last_update.timestamp()),
+                "centres_by_date": {
+                    "20210306": ["hosp1", "hosp2"],
+                    "20210509": ["hosp3", "hosp4"]
+                }
+            })
+        }
+    }
+
+    centres_by_date, received_last_update = db.get_min_date_info()
 
+    assert centres_by_date == {
+        datetime(2021, 3, 6): ["hosp1", "hosp2"],
+        datetime(2021, 5, 9): ["hosp3", "hosp4"]
+    }
+    assert received_last_update == datetime.fromtimestamp(int(last_update.timestamp()))