From a1965c38aff72e0477daf70444d23ac28ebb26e6 Mon Sep 17 00:00:00 2001
From: SmCTwelve <smctwelve@gmail.com>
Date: Sun, 17 Sep 2023 03:00:52 +0100
Subject: [PATCH] Add fallback for missing ergast drivers in round

---
 .vscode/settings.json            |  2 --
 f1/api/ergast.py                 | 20 +++++++++----
 f1/tests/mock_response/models.py | 48 ++++++++++++++++++--------------
 f1/tests/test_data.py            | 17 +++++++++++
 4 files changed, 59 insertions(+), 28 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 4e081c4..38dc48a 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -3,8 +3,6 @@
   "python.testing.pytestEnabled": false,
   "python.testing.unittestEnabled": true,
   "editor.formatOnSave": true,
-  "python.formatting.provider": "autopep8",
-  "python.formatting.autopep8Args": ["--max-line-length", "120", "--experimental"],
   "[python]": {},
   "python.analysis.typeCheckingMode": "off"
 }
diff --git a/f1/api/ergast.py b/f1/api/ergast.py
index 6d42d6c..f40a3ce 100644
--- a/f1/api/ergast.py
+++ b/f1/api/ergast.py
@@ -72,16 +72,26 @@ async def get_all_drivers(season=None, round=None) -> list[dict]:
 
     Raises `MissingDataError`.
     """
+    round_url = f'{BASE_URL}/{season}/{round}/drivers.json?limit=1000'
+    season_url = f'{BASE_URL}/{season}/drivers.json?limit=1000'
+    all_time_url = f'{BASE_URL}/drivers.json?limit=1000'
+
+    # Get JSON data as dict
     if season and round:
-        url = f'{BASE_URL}/{season}/{round}/drivers.json?limit=1000'
+        res = await fetch(round_url)
     elif season:
-        url = f'{BASE_URL}/{season}/drivers.json?limit=1000'
+        res = await fetch(season_url)
     else:
-        url = f'{BASE_URL}/drivers.json?limit=1000'
-    # Get JSON data as dict
-    res = await fetch(url)
+        res = await fetch(all_time_url)
+
     if res is None:
         raise MissingDataError()
+
+    # Fallback to season only if data for the round is unavailable
+    if round and not res['MRData']['DriverTable']['Drivers']:
+        logger.warning("Driver data for specified round is missing, falling back to season list.")
+        res = await fetch(season_url)
+
     return res['MRData']['DriverTable']['Drivers']
 
 
diff --git a/f1/tests/mock_response/models.py b/f1/tests/mock_response/models.py
index 5ba5606..81a664f 100644
--- a/f1/tests/mock_response/models.py
+++ b/f1/tests/mock_response/models.py
@@ -386,25 +386,31 @@ def result_wrapper(body, quali=False):
     ]
 }
 
-all_drivers = [
-    {
-        "driverId": "alonso",
-        "permanentNumber": "14",
-        "code": "ALO",
-        "url": "http:\/\/en.wikipedia.org\/wiki\/Fernando_Alonso",
-        "givenName": "Fernando",
-        "familyName": "Alonso",
-        "dateOfBirth": "1981-07-29",
-        "nationality": "Spanish"
-    },
-    {
-        "driverId": "max_verstappen",
-        "permanentNumber": "1",
-        "code": "VER",
-        "url": "http:\/\/en.wikipedia.org\/wiki\/Max_Verstappen",
-        "givenName": "Max",
-        "familyName": "Verstappen",
-        "dateOfBirth": "1981-07-29",
-        "nationality": "Dutch"
+all_drivers = {
+    "MRData": {
+        "DriverTable": {
+            "Drivers": [
+                {
+                    "driverId": "alonso",
+                    "permanentNumber": "14",
+                    "code": "ALO",
+                    "url": "http:\/\/en.wikipedia.org\/wiki\/Fernando_Alonso",
+                    "givenName": "Fernando",
+                    "familyName": "Alonso",
+                    "dateOfBirth": "1981-07-29",
+                    "nationality": "Spanish"
+                },
+                {
+                    "driverId": "max_verstappen",
+                    "permanentNumber": "1",
+                    "code": "VER",
+                    "url": "http:\/\/en.wikipedia.org\/wiki\/Max_Verstappen",
+                    "givenName": "Max",
+                    "familyName": "Verstappen",
+                    "dateOfBirth": "1981-07-29",
+                    "nationality": "Dutch"
+                }
+            ]
+        }
     }
-]
+}
diff --git a/f1/tests/test_data.py b/f1/tests/test_data.py
index be25d02..ae4ffc8 100644
--- a/f1/tests/test_data.py
+++ b/f1/tests/test_data.py
@@ -444,6 +444,23 @@ async def test_get_driver_career(self, mock_fetch):
         self.check_total_and_num_results(data['Seasons']['total'], data['Seasons']['years'])
         self.check_total_and_num_results(data['Teams']['total'], data['Teams']['names'])
 
+    @patch(fetch_path)
+    @async_test
+    async def test_get_all_drivers(self, mock_fetch):
+        mock_fetch.return_value = models.all_drivers
+        res = await ergast.get_all_drivers()
+        self.assertEqual(res[0]["code"], "ALO")
+        self.assertEqual(res[1]["code"], "VER")
+
+    @patch(fetch_path)
+    @async_test
+    async def test_get_all_drivers_missing_round(self, mock_fetch):
+        empty_data = {"MRData": {"DriverTable": {"Drivers": []}}}
+        mock_fetch.side_effect = [empty_data, models.all_drivers]
+        res = await ergast.get_all_drivers(round=100)
+        self.assertEqual(mock_fetch.call_count, 2)
+        self.assertEqual(len(res), 2)
+
     # boundary tests