Wrap responses in a new type to include credits and throttling data

laurentpellegrino · laurentpellegrino · commit ca53a7e672cb · 2024-06-13T11:27:14.000+02:00
This fixes #44.
diff --git a/ipregistry/core.py b/ipregistry/core.py
@@ -15,7 +15,7 @@
 """
 
 from .cache import IpregistryCache, NoCache
-from .model import LookupError
+from .model import LookupError, ApiResponse, ApiResponseCredits, ApiResponseThrottling
 from .request import DefaultRequestHandler, IpregistryRequestHandler
 
 
@@ -30,7 +30,7 @@ def __init__(self, key_or_config, **kwargs):
         if not isinstance(self._requestHandler, IpregistryRequestHandler):
             raise ValueError("Given request handler instance is not of type IpregistryRequestHandler")
 
-    def batch_lookup_ips(self, ips, options):
+    def batch_lookup_ips(self, ips, **options):
         sparse_cache = [None] * len(ips)
         cache_misses = []
 
@@ -44,21 +44,32 @@ def batch_lookup_ips(self, ips, options):
                 sparse_cache[i] = cache_value
 
         result = [None] * len(ips)
-        fresh_ip_info = self._requestHandler.batch_lookup_ips(cache_misses, options)
+        if len(cache_misses) > 0:
+            response = self._requestHandler.batch_lookup_ips(cache_misses, options)
+        else:
+            response = ApiResponse(
+                ApiResponseCredits(),
+                [],
+                ApiResponseThrottling()
+            )
+
+        fresh_ip_info = response.data
         j = 0
         k = 0
 
-        for cachedIpInfo in sparse_cache:
-            if cachedIpInfo is None:
+        for cached_ip_info in sparse_cache:
+            if cached_ip_info is None:
                 if not isinstance(fresh_ip_info[k], LookupError):
                     self._cache.put(self.__build_cache_key(ips[j], options), fresh_ip_info[k])
                 result[j] = fresh_ip_info[k]
                 k += 1
             else:
-                result[j] = cachedIpInfo
+                result[j] = cached_ip_info
             j += 1
 
-        return result
+        response.data = result
+
+        return response
 
     def lookup_ip(self, ip='', **options):
         if isinstance(ip, str) and len(ip) > 0:
@@ -74,10 +85,15 @@ def __lookup_ip(self, ip, options):
         cache_value = self._cache.get(cache_key)
 
         if cache_value is None:
-            cache_value = self._requestHandler.lookup_ip(ip, options)
-            self._cache.put(cache_key, cache_value)
-
-        return cache_value
+            response = self._requestHandler.lookup_ip(ip, options)
+            self._cache.put(cache_key, response.data)
+            return response
+
+        return ApiResponse(
+            ApiResponseCredits(),
+            cache_value,
+            ApiResponseThrottling()
+        )
 
     @staticmethod
     def __build_cache_key(ip, options):
diff --git a/ipregistry/model.py b/ipregistry/model.py
@@ -15,6 +15,41 @@
 """
 
 from .json import *
+from typing import Generic, TypeVar, Optional, Dict, Any
+
+T = TypeVar('T')
+
+
+class ApiResponseCredits:
+    def __init__(self, consumed: Optional[int] = 0, remaining: Optional[int] = None):
+        self.consumed = consumed
+        self.remaining = remaining
+
+    def __str__(self):
+        fields = ', '.join(f"{key}={value}" for key, value in self.__dict__.items())
+        return f"{self.__class__.__name__}({fields})"
+
+
+class ApiResponseThrottling:
+    def __init__(self, limit: int = None, remaining: int = None, reset: int = None):
+        self.limit = limit
+        self.remaining = remaining
+        self.reset = reset
+
+    def __str__(self):
+        fields = ', '.join(f"{key}={value}" for key, value in self.__dict__.items())
+        return f"{self.__class__.__name__}({fields})"
+
+
+class ApiResponse(Generic[T]):
+    def __init__(self, credits: ApiResponseCredits, data: T, throttling: Optional[ApiResponseThrottling] = None):
+        self.credits = credits
+        self.data = data
+        self.throttling = throttling
+
+    def __str__(self):
+        fields = ', '.join(f"{key}={value}" for key, value in self.__dict__.items())
+        return f"{self.__class__.__name__}({fields})"
 
 
 class RequesterIpInfo(IpInfo):
@@ -26,12 +61,20 @@ class IpregistryError(Exception):
 
 
 class ApiError(IpregistryError):
-    pass
+    def __init__(self, code: str, message: str, resolution: str):
+        self.code = code
+        self.message = message
+        self.resolution = resolution
 
 
 class ClientError(IpregistryError):
    pass
 
 
-class LookupError(IpregistryError):
-    pass
+class LookupError(ApiError):
+    def __init__(self, json: Dict[str, Any]):
+        super().__init__(
+            code=json.get('code'),
+            message=json.get('message'),
+            resolution=json.get('resolution')
+        )
diff --git a/ipregistry/request.py b/ipregistry/request.py
@@ -14,14 +14,16 @@
     limitations under the License.
 """
 
-from abc import ABC, abstractmethod
 import json
-import requests
 import sys
 import urllib.parse
+from abc import ABC, abstractmethod
+from typing import Union
+
+import requests
 
 from .__init__ import __version__
-from .model import ApiError, ClientError, IpInfo
+from .model import ApiError, ApiResponse, ApiResponseCredits, ApiResponseThrottling, ClientError, IpInfo, LookupError
 
 
 class IpregistryRequestHandler(ABC):
@@ -53,28 +55,86 @@ def _build_base_url(self, ip, options):
 
 class DefaultRequestHandler(IpregistryRequestHandler):
     def batch_lookup_ips(self, ips, options):
+        response = None
         try:
-            r = requests.post(self._build_base_url('', options), data=json.dumps(ips), headers=self.__headers(), timeout=self._config.timeout)
-            r.raise_for_status()
-            return list(map(lambda data: LookupError(data) if 'code' in data else IpInfo(**data), r.json()['results']))
+            response = requests.post(
+                self._build_base_url('', options),
+                data=json.dumps(ips),
+                headers=self.__headers(),
+                timeout=self._config.timeout
+            )
+            response.raise_for_status()
+            results = response.json().get('results', [])
+
+            parsed_results = [
+                LookupError(data) if 'code' in data else IpInfo(**data)
+                for data in results
+            ]
+
+            return self.build_api_response(response, parsed_results)
         except requests.HTTPError:
-            raise ApiError(r.json())
+            self.__create_api_error(response)
         except Exception as e:
             raise ClientError(e)
 
     def lookup_ip(self, ip, options):
+        response = None
         try:
-            r = requests.get(self._build_base_url(ip, options), headers=self.__headers(), timeout=self._config.timeout)
-            r.raise_for_status()
-            return IpInfo(**r.json())
+            response = requests.get(
+                self._build_base_url(ip, options),
+                headers=self.__headers(),
+                timeout=self._config.timeout
+            )
+            response.raise_for_status()
+            json_response = response.json()
+
+            return self.build_api_response(response, IpInfo(**json_response))
         except requests.HTTPError:
-            raise ApiError(r.json())
-        except Exception as e:
-            raise ClientError(e)
+            self.__create_api_error(response)
+        except Exception as err:
+            raise ClientError(err)
 
     def origin_lookup_ip(self, options):
         return self.lookup_ip('', options)
 
+    @staticmethod
+    def build_api_response(response, data):
+        throttling_limit = DefaultRequestHandler.__convert_to_int(response.headers.get('x-rate-limit-limit'))
+        throttling_remaining = DefaultRequestHandler.__convert_to_int(response.headers.get('x-rate-limit-remaining'))
+        throttling_reset = DefaultRequestHandler.__convert_to_int(response.headers.get('x-rate-limit-reset'))
+
+        ipregistry_credits_consumed = DefaultRequestHandler.__convert_to_int(response.headers.get('ipregistry-credits-consumed'))
+        ipregistry_credits_remaining = DefaultRequestHandler.__convert_to_int(response.headers.get('ipregistry-credits-remaining'))
+
+        return ApiResponse(
+            ApiResponseCredits(
+                ipregistry_credits_consumed,
+                ipregistry_credits_remaining,
+            ),
+            data,
+            None if throttling_limit is None and throttling_remaining is None and throttling_reset is None else
+            ApiResponseThrottling(
+                throttling_limit,
+                throttling_remaining,
+                throttling_reset
+            )
+        )
+
+    @staticmethod
+    def __convert_to_int(value: str) -> Union[int, None]:
+        try:
+            return int(value)
+        except (ValueError, TypeError):
+            return None
+
+    @staticmethod
+    def __create_api_error(response):
+        if response is not None:
+            json_response = response.json()
+            raise ApiError(json_response['code'], json_response['message'], json_response['resolution'])
+        else:
+            raise ClientError("HTTP Error occurred, but no response was received.")
+
     @staticmethod
     def __headers():
         return {
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -17,37 +17,109 @@
 import os
 import unittest
 
+from ipregistry import ApiError, IpInfo, LookupError, ClientError
 from ipregistry.cache import InMemoryCache, NoCache
-from ipregistry.core import IpregistryClient
+from ipregistry.core import IpregistryClient, IpregistryConfig
 
 
 class TestIpregistryClient(unittest.TestCase):
-    def test_defaultclient_cache(self):
+
+    def test_batch_lookup_ips(self):
+        """
+        Test batch ips lookup with valid and invalid inputs
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'))
+        response = client.batch_lookup_ips(['1.1.1.1', 'invalid', '8.8.8.8'])
+        self.assertEqual(3, len(response.data))
+        self.assertEqual(True, isinstance(response.data[0], IpInfo))
+        self.assertEqual(True, isinstance(response.data[1], LookupError))
+        self.assertEqual('INVALID_IP_ADDRESS', response.data[1].code)
+        self.assertEqual(True, isinstance(response.data[2], IpInfo))
+
+    def test_client_cache_default(self):
         """
         Test that default cache is an instance of NoCache
         """
-        client = IpregistryClient("")
+        client = IpregistryClient("tryout")
         self.assertEqual(True, isinstance(client._cache, NoCache))
 
-    def test_simple_lookup(self):
+    def test_client_cache_inmemory_ip_lookup(self):
+        """
+        Test the client in memory cache
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'), cache=InMemoryCache(maxsize=2048, ttl=600))
+        lookup_ip_response = client.lookup_ip('1.1.1.3')
+        lookup_ip_response2 = client.lookup_ip('1.1.1.3')
+
+        self.assertEqual(1, lookup_ip_response.credits.consumed)
+        self.assertEqual(0, lookup_ip_response2.credits.consumed)
+
+    def test_client_cache_inmemory_batch_ips_lookup(self):
+        """
+        Test the client in memory cache
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'), cache=InMemoryCache(maxsize=2048, ttl=600))
+        lookup_ip_response = client.lookup_ip('1.1.1.3')
+        batch_ips_response = client.batch_lookup_ips(['1.1.1.1', '1.1.1.3'])
+
+        self.assertEqual(1, lookup_ip_response.credits.consumed)
+        self.assertEqual(1, batch_ips_response.credits.consumed)
+
+        batch_ips_response2 = client.batch_lookup_ips(['1.1.1.1', '1.1.1.3'])
+        self.assertEqual(0, batch_ips_response2.credits.consumed)
+
+    def test_lookup_ip(self):
         """
         Test that a simple lookup returns data
         """
         client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'))
         response = client.lookup_ip('8.8.8.8')
-        self.assertIsNotNone(response.ip)
-        self.assertIsNotNone(response.company.domain)
-        self.assertEqual('US', response.location.country.code)
+        self.assertIsNotNone(response.data.ip)
+        self.assertIsNotNone(response.data.company.domain)
+        self.assertEqual('US', response.data.location.country.code)
 
-    def test_simple_lookup_in_memory_cache(self):
+    def test_lookup_ip_invalid_input(self):
+        """
+        Test that an IP lookup with an invalid input fails with an ApiError
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'))
+        with self.assertRaises(ApiError) as context:
+            response = client.lookup_ip('invalid')
+        print("test", context.exception)
+        self.assertEqual('INVALID_IP_ADDRESS', context.exception.code)
+
+    def test_lookup_ip_cache(self):
         """
         Test consecutive lookup with in-memory cache
         """
         client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'), cache=InMemoryCache(maxsize=2048, ttl=600))
         response = client.lookup_ip('8.8.8.8')
         response = client.lookup_ip('8.8.8.8')
-        self.assertIsNotNone(response.ip)
-        self.assertIsNotNone(response.company.domain)
+        self.assertIsNotNone(response.data.ip)
+        self.assertIsNotNone(response.data.company.domain)
+
+    def test_response_metadata(self):
+        """
+        Test metadata returned for each successful response
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'))
+        batch_ips_lookup_response = client.batch_lookup_ips(['1.1.1.1', 'invalid', '8.8.8.8'])
+        lookup_ip_response = client.lookup_ip('1.1.1.2')
+
+        self.assertEqual(3, batch_ips_lookup_response.credits.consumed)
+        self.assertEqual(1, lookup_ip_response.credits.consumed)
+
+        self.assertIsNotNone(batch_ips_lookup_response.credits.remaining)
+        self.assertIsNotNone(lookup_ip_response.credits.remaining)
+
+    def test_lookup_timeout(self):
+        """
+        Test a client error is raised upon connection timeout
+        """
+        client = IpregistryClient(IpregistryConfig(os.getenv('IPREGISTRY_API_KEY'), "https://api.ipregistry.co",
+                                                   0.0001))
+        with self.assertRaises(ClientError):
+            client.lookup_ip('1.1.1.1')
 
 
 if __name__ == '__main__':
