From 1c652ebec3959e0f6bd2fd21e90fcca8432f43f8 Mon Sep 17 00:00:00 2001
From: Appelsiini1 <37345715+Appelsiini1@users.noreply.github.com>
Date: Sun, 23 Jan 2022 14:15:46 +0200
Subject: [PATCH] Add error handling for HTTP requests in vaccine command

---
 constants.py         |  2 +-
 modules/vaccine.py   | 51 +++++++++++++++++++++++++++-----------------
 webhook/versions.txt |  3 +++
 3 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/constants.py b/constants.py
index dec1cfc..5536caa 100644
--- a/constants.py
+++ b/constants.py
@@ -135,4 +135,4 @@ def resetLimits(self):
 REQUEST_LIMITS = requestLimits()
 
 # Version number
-VERSION = "v2.0.3"
+VERSION = "v2.0.4"
diff --git a/modules/vaccine.py b/modules/vaccine.py
index a103e0b..bb12283 100644
--- a/modules/vaccine.py
+++ b/modules/vaccine.py
@@ -117,24 +117,35 @@ class Vaccine(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
 
-    def getVaccineInfo(self, param):
-        response = requests.get(
-            f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-533170&row=area-{param}",
-            headers={"User-Agent": "Appelsiini1:n Discord Botti"},
-        )
-        response2 = requests.get(
-            f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-533164&row=area-{param}",
-            headers={"User-Agent": "Appelsiini1:n Discord Botti"},
-        )
-        response3 = requests.post(
-            "https://pxnet2.stat.fi:443/PXWeb/api/v1/fi/StatFin/vrm/vaerak/statfin_vaerak_pxt_11ra.px",
-            headers={"User-Agent": "Appelsiini1:n Discord Botti"},
-            json=QUERY,
-        )
-        response4 = requests.post(
-            f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-639082&row=area-{param}",
-            headers={"User-Agent": "Appelsiini1:n Discord Botti"},
-        )
+    def getVaccineInfo(self, param, ctx):
+        try:
+            response = requests.get(
+                f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-533170&row=area-{param}",
+                headers={"User-Agent": "Appelsiini1:n Discord Botti"},
+            )
+            response2 = requests.get(
+                f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-533164&row=area-{param}",
+                headers={"User-Agent": "Appelsiini1:n Discord Botti"},
+            )
+            response3 = requests.post(
+                "https://pxnet2.stat.fi:443/PXWeb/api/v1/fi/StatFin/vrm/vaerak/statfin_vaerak_pxt_11ra.px",
+                headers={"User-Agent": "Appelsiini1:n Discord Botti"},
+                json=QUERY,
+            )
+            response4 = requests.post(
+                f"https://sampo.thl.fi/pivot/prod/fi/vaccreg/cov19cov/fact_cov19cov.json?row=cov_vac_dose-639082&row=area-{param}",
+                headers={"User-Agent": "Appelsiini1:n Discord Botti"},
+            )
+        except requests.exceptions.ConnectionError:
+            emb = discord.Embed(color=get_hex_colour(error=True))
+            emb.description = "Could not fetch vaccination data from THL or population data from Statistics Finland.\n[Connection timed out]"
+            ctx.send(embed=emb)
+            return
+        except Exception:
+            emb = discord.Embed(color=get_hex_colour(error=True))
+            emb.description = "Could not fetch vaccination data from THL or population data from Statistics Finland.\n[General error]"
+            ctx.send(embed=emb)
+            return
         if (
             response.status_code == 200
             and response2.status_code == 200
@@ -212,7 +223,7 @@ async def sendVaccInfo(self, ctx):
             await forbiddenErrorHandler(ctx.message)
         
         if len(ctx.message.content.split(" ")) == 2:
-            result = self.getVaccineInfo("518362")
+            result = self.getVaccineInfo("518362", ctx)
             if type(result) == VACC_DATA:
                 emb = self.makeEmbed(emb, vaccData=result)
             else:
@@ -240,7 +251,7 @@ async def sendVaccInfo(self, ctx):
                     await s_msg.edit(embed=emb)
                     return
                 
-                result = self.getVaccineInfo(param)
+                result = self.getVaccineInfo(param, ctx)
                 if type(result) == VACC_DATA:
                     emb = self.makeEmbed(emb, vaccData=result)
                 else:
diff --git a/webhook/versions.txt b/webhook/versions.txt
index 92da08d..7754d0e 100644
--- a/webhook/versions.txt
+++ b/webhook/versions.txt
@@ -1,3 +1,6 @@
+v2.0.4
+- Add error handling for HTTP requests in vaccine command
+###
 v2.0.3
 - Fix area name in vaccine command
 - Rework `!c mood` command