From 572b177b665d69775b02b64ca3f09f212398b345 Mon Sep 17 00:00:00 2001
From: Roel de Vries <roel.de.vries@amsterdam.nl>
Date: Wed, 24 Jan 2024 14:28:22 +0100
Subject: [PATCH] Retry on timeouts during getting remote amschema

---
 gobcore/model/amschema/repo.py            | 18 ++++++++++++++----
 tests/gobcore/model/amschema/test_repo.py | 21 +++++++++++++++++----
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/gobcore/model/amschema/repo.py b/gobcore/model/amschema/repo.py
index f4433a8f..0e60ff54 100644
--- a/gobcore/model/amschema/repo.py
+++ b/gobcore/model/amschema/repo.py
@@ -2,8 +2,9 @@
 
 import os
 
-import requests
-from pydash import snake_case
+from requests import Session
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
 
 from gobcore.model import Schema
 from gobcore.model.amschema.model import Dataset, Table
@@ -25,8 +26,17 @@ def _get_file(self, location: str):
             return self._load_file(location)
 
     def _download_file(self, location: str):
-        r = requests.get(location, timeout=5)
-        r.raise_for_status()
+        retries = Retry(
+            total=6,
+            backoff_factor=1,  # 1 x 2^0 until 1 x 2^6
+            status_forcelist=[502, 503, 504],
+            allowed_methods={"GET"},
+        )
+
+        with Session() as session:
+            session.mount("https://", HTTPAdapter(max_retries=retries))
+            r = session.get(location, timeout=10)
+            r.raise_for_status()
 
         return r.json()
 
diff --git a/tests/gobcore/model/amschema/test_repo.py b/tests/gobcore/model/amschema/test_repo.py
index 0a756f7f..23b48ed7 100644
--- a/tests/gobcore/model/amschema/test_repo.py
+++ b/tests/gobcore/model/amschema/test_repo.py
@@ -57,12 +57,16 @@ def test_repo_base(self):
             instance.get_schema(schema)
             instance._download_dataset.assert_called_with("https://test.repo/datasets/nap/dataset.json")
 
-    @patch("gobcore.model.amschema.repo.requests")
-    def test_download_dataset(self, mock_requests):
+    @patch("gobcore.model.amschema.repo.HTTPAdapter")
+    @patch("gobcore.model.amschema.repo.Retry")
+    @patch("gobcore.model.amschema.repo.Session")
+    def test_download_dataset(self, mock_session, mock_retry, mock_http):
         """Test remote (HTTP) AMS schema dataset."""
+        mock_session_cm = mock_session.return_value.__enter__.return_value
+
         with open(Path(__file__).parent.parent.parent.joinpath('amschema_fixtures/dataset.json')) as f:
             filecontents = f.read()
-            mock_requests.get.return_value.json = lambda: json.loads(filecontents)
+            mock_session_cm.get.return_value.json = lambda: json.loads(filecontents)
 
         instance = AMSchemaRepository()
         dataset = get_dataset()
@@ -71,7 +75,16 @@ def test_download_dataset(self, mock_requests):
         result = instance._download_dataset(download_url)
         self.assertEqual(result, dataset)
 
-        mock_requests.get.assert_called_with(download_url, timeout=5)
+        mock_retry.assert_called_with(
+            total=6,
+            backoff_factor=1,
+            status_forcelist=[502, 503, 504],
+            allowed_methods={"GET"}
+        )
+        mock_http.assert_called_with(max_retries=mock_retry.return_value)
+        mock_session_cm.get.assert_called_with(download_url, timeout=10)
+        mock_session_cm.get.return_value.raise_for_status.assert_called_once()
+        mock_session_cm.mount.assert_called_with("https://", mock_http.return_value)
 
     @patch("gobcore.model.amschema.repo.json_to_cached_dict")
     def test_local_table(self, mock_cached_dict):