From 7731258f2f1c189bf92a7d8a51eba0eea85169d8 Mon Sep 17 00:00:00 2001
From: Philip Meier <github.pmeier@posteo.de>
Date: Thu, 24 Sep 2020 11:16:55 +0200
Subject: [PATCH] limit requests per time in download tests (#2699)

---
 test/test_datasets_download.py | 35 +++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/test/test_datasets_download.py b/test/test_datasets_download.py
index 6fe86ce1449..7b189368dda 100644
--- a/test/test_datasets_download.py
+++ b/test/test_datasets_download.py
@@ -1,9 +1,11 @@
 import contextlib
 import itertools
+import time
 import unittest
 import unittest.mock
+from datetime import datetime
 from os import path
-from time import sleep
+from urllib.parse import urlparse
 from urllib.request import urlopen, Request
 
 from torchvision import datasets
@@ -13,6 +15,34 @@
 from fakedata_generation import places365_root
 
 
+def limit_requests_per_time(min_secs_between_requests=2.0):
+    last_requests = {}
+
+    def outer_wrapper(fn):
+        def inner_wrapper(request, *args, **kwargs):
+            url = request.full_url if isinstance(request, Request) else request
+
+            netloc = urlparse(url).netloc
+            last_request = last_requests.get(netloc)
+            if last_request is not None:
+                elapsed_secs = (datetime.now() - last_request).total_seconds()
+                delta = min_secs_between_requests - elapsed_secs
+                if delta > 0:
+                    time.sleep(delta)
+
+            response = fn(request, *args, **kwargs)
+            last_requests[netloc] = datetime.now()
+
+            return response
+
+        return inner_wrapper
+
+    return outer_wrapper
+
+
+urlopen = limit_requests_per_time()(urlopen)
+
+
 class DownloadTester(unittest.TestCase):
     @staticmethod
     @contextlib.contextmanager
@@ -37,7 +67,7 @@ def retry(fn, times=1, wait=5.0):
                 return fn()
             except AssertionError as error:
                 msgs.append(str(error))
-                sleep(wait)
+                time.sleep(wait)
         else:
             raise AssertionError(
                 "\n".join(
@@ -80,7 +110,6 @@ def test_download(self):
         for url, md5 in self.collect_urls_and_md5s():
             with self.subTest(url=url, md5=md5):
                 self.retry(lambda: assert_fn(url, md5))
-                sleep(2.0)
 
     def collect_urls_and_md5s(self):
         raise NotImplementedError