updating handler function

Author: Samuel-Gill

ipinfo/handler.py

@@ -315,8 +315,6 @@ def getIterativeBatchDetails(
         if batch_size is None:
             batch_size = BATCH_MAX_SIZE
 
-        results = {}
-
         lookup_addresses = []
         for ip_address in ip_addresses:
             if isinstance(ip_address, IPv4Address) or isinstance(
@@ -329,16 +327,9 @@ def getIterativeBatchDetails(
                 details["ip"] = ip_address
                 details["bogon"] = True
                 yield Details(details)
-
-            try:
-                cached_ipaddr = self.cache[cache_key(ip_address)]
-                results[ip_address] = cached_ipaddr
-            except KeyError:
+            else:
                 lookup_addresses.append(ip_address)
 
-        if len(lookup_addresses) == 0:
-            yield from results.items()
-
         url = API_URL + "/batch"
         headers = handler_utils.get_headers(self.access_token, self.headers)
         headers["content-type"] = "application/json"
@@ -348,30 +339,28 @@ def getIterativeBatchDetails(
             try:
                 response = requests.post(url, json=batch, headers=headers)
             except Exception as e:
-                return handler_utils.return_or_fail(raise_on_fail, e, results)
+                return handler_utils.return_or_fail(raise_on_fail, e)
 
             try:
                 if response.status_code == 429:
                     raise RequestQuotaExceededError()
                 response.raise_for_status()
             except Exception as e:
-                return handler_utils.return_or_fail(raise_on_fail, e, results)
-
-            json_response = response.json()
-            for ip_address, details in json_response.items():
-                self.cache[cache_key(ip_address)] = details
-
-            results.update(json_response)
-
-            for detail in results.values():
-                if isinstance(detail, dict):
-                    handler_utils.format_details(
-                        detail,
-                        self.countries,
-                        self.eu_countries,
-                        self.countries_flags,
-                        self.countries_currencies,
-                        self.continents,
-                    )
-
-            yield from results.items()
+                return handler_utils.return_or_fail(raise_on_fail, e)
+
+            details = response.json()
+
+            # format & cache
+            handler_utils.format_details(
+                details,
+                self.countries,
+                self.eu_countries,
+                self.countries_flags,
+                self.countries_currencies,
+                self.continents,
+            )
+            for ip in batch:
+                detail = details.get(ip)
+                if detail is not None:
+                    self.cache[cache_key(ip)] = detail
+                    yield detail

ipinfo/handler_utils.py

@@ -80,7 +80,7 @@ def format_details(
     details["country_name"] = countries.get(details.get("country"))
     details["isEU"] = details.get("country") in eu_countries
     details["country_flag_url"] = (
-        COUNTRY_FLAGS_URL + details.get("country") + ".svg"
+        COUNTRY_FLAGS_URL + (details.get("country") or "") + ".svg"
     )
     details["country_flag"] = copy.deepcopy(
         countries_flags.get(details.get("country"))

tests/handler_test.py

@@ -108,7 +108,7 @@ def test_get_details():
 
 def _prepare_batch_test():
     """Helper for preparing batch test cases."""
-    token = os.environ.get("IPINFO_TOKEN", "")
+    token =  os.environ.get("IPINFO_TOKEN", "")
     if not token:
         pytest.skip("token required for batch tests")
     handler = Handler(token)
@@ -131,17 +131,17 @@ def _check_batch_details(ips, details, token):
             assert "domains" in d
 
 
-def _check_iterative_batch_details(ip, details, token):
+def _check_iterative_batch_details(details, token):
     """Helper for iterative batch tests."""
-    assert ip == details.get("ip")
-    assert "country" in details
-    assert "city" in details
+    assert "ip" in details, "Key 'ip' not found in details"
+    assert "country" in details, "Key 'country' not found in details"
+    assert "city" in details, "Key 'city' not found in details"
     if token:
-        assert "asn" in details
-        assert "company" in details
-        assert "privacy" in details
-        assert "abuse" in details
-        assert "domains" in details
+        assert "asn" in details, "Key 'asn' not found in details"
+        assert "company" in details, "Key 'company' not found in details"
+        assert "privacy" in details, "Key 'privacy' not found in details"
+        assert "abuse" in details, "Key 'abuse' not found in details"
+        assert "domains" in details, "Key 'domains' not found in details"
 
 
 @pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
@@ -166,8 +166,8 @@ def test_get_iterative_batch_details(batch_size):
     details_iterator = handler.getIterativeBatchDetails(
         ips, batch_size=batch_size
     )
-    for ip, details in details_iterator:
-        _check_iterative_batch_details(ip, details, token)
+    for details in details_iterator:
+        _check_iterative_batch_details(details, token)
 
 
 #############