From 5ad0c9adcfb1f5fd294ce1525ab18ebe1215ccfc Mon Sep 17 00:00:00 2001
From: "romain.fontugne" <romain.fontugne@gmail.com>
Date: Mon, 13 Aug 2018 13:01:21 +0900
Subject: [PATCH] Allow multiple measurement IDs for AtlasResultsRequest and
 fetch results asynchronously

---
 ripe/atlas/cousteau/request.py | 56 ++++++++++++++++++++++++++++------
 setup.py                       |  3 +-
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/ripe/atlas/cousteau/request.py b/ripe/atlas/cousteau/request.py
index 1c8ce1b..6b76ded 100644
--- a/ripe/atlas/cousteau/request.py
+++ b/ripe/atlas/cousteau/request.py
@@ -19,6 +19,7 @@
 """
 
 import calendar
+import grequests 
 import requests
 from dateutil import parser
 from datetime import datetime
@@ -38,6 +39,12 @@ class AtlasRequest(object):
         "DELETE": requests.delete
     }
 
+    http_methods_async = {
+        "GET": grequests.get,
+        "POST": grequests.post,
+        "DELETE": grequests.delete
+    }
+
     def __init__(self, **kwargs):
 
         self.url = ""
@@ -72,6 +79,10 @@ def get_headers(self):
 
         return headers
 
+    def exception_handler(self, request, exc):
+        """Exception handler for asynchronous results fetching."""
+        return exc.args
+
     def http_method(self, method):
         """
         Execute the given HTTP method and returns if it's success or not
@@ -82,12 +93,27 @@ def http_method(self, method):
 
         try:
             response = self.get_http_method(method)
-            is_success = response.ok
-
-            try:
-                response_message = response.json()
-            except ValueError:
-                response_message = response.text
+            if isinstance(response, list):
+                is_success = []
+                response_message = []
+
+                for resp in response:
+                    if isinstance(resp, requests.Response):
+                        is_success.append(resp.ok)
+                        try:
+                            response_message.append(resp.json())
+                        except ValueError:
+                            response_message.append(resp.text)
+                    else:
+                        # An exception occured
+                        is_success.append(False)
+                        response_message.append(resp)
+            else:
+                is_success = response.ok
+                try:
+                    response_message = response.json()
+                except ValueError:
+                    response_message = response.text
 
         except requests.exceptions.RequestException as exc:
             is_success = False
@@ -97,14 +123,23 @@ def http_method(self, method):
 
     def get_http_method(self, method):
         """Gets the http method that will be called from the requests library"""
-        return self.http_methods[method](self.url, **self.http_method_args)
+        if isinstance(self.url, list):
+            rs = [self.http_methods_async[method](url, **self.http_method_args) 
+                    for url in self.url]
+            return grequests.map(rs, exception_handler=self.exception_handler)
+        else:
+            return self.http_methods[method](self.url, **self.http_method_args)
 
     def build_url(self):
         """
         Builds the request's url combining server and url_path
         classes attributes.
         """
-        self.url = "https://{0}{1}".format(self.server, self.url_path)
+        if isinstance(self.url_path, list):
+            self.url = ["https://{0}{1}".format(self.server, url) 
+                    for url in self.url_path]
+        else:
+            self.url = "https://{0}{1}".format(self.server, self.url_path)
 
     def get(self, **url_params):
         """
@@ -312,7 +347,10 @@ def __init__(self, **kwargs):
 
         self.probe_ids = self.clean_probes(kwargs.get("probe_ids"))
 
-        self.url_path = self.url_path.format(self.msm_id)
+        if isinstance(self.msm_id, list):
+            self.url_path = [self.url_path.format(mid) for mid in self.msm_id]
+        else:
+            self.url_path = self.url_path.format(self.msm_id)
 
         self.update_http_method_params()
 
diff --git a/setup.py b/setup.py
index 360382d..cb6b009 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,8 @@
 
 install_requires = [
     "python-dateutil", "socketIO-client>=0.6.5",
-    "requests>=2.7.0", "websocket-client<0.99"
+    "requests>=2.7.0", "websocket-client<0.99",
+    "grequests"
 ]
 
 tests_require = [