Fix batch lookups with a fields selection and fail cleanly on short batch responses

Author: laurentpellegrino

ipregistry/async_client.py

@@ -24,7 +24,8 @@
                     ClientError, IpregistryLookupError, RequesterAutonomousSystem,
                     RequesterIpInfo, RequesterUserAgent, UserAgent)
 from .request import (DefaultRequestHandler, IpregistryRequestHandler, _backoff_interval,
-                      _build_headers, _is_retryable_status, _raise_api_error, _retry_delay)
+                      _build_headers, _is_retryable_status, _prefix_batch_fields,
+                      _raise_api_error, _retry_delay)
 
 try:
     import httpx
@@ -57,7 +58,7 @@ async def aclose(self):
     async def batch_lookup_asns(self, asns, options):
         response = await self._request_with_retry(
             'POST',
-            self._build_base_url('', options),
+            self._build_base_url('', _prefix_batch_fields(options)),
             data=json.dumps(["AS" + str(asn) for asn in asns])
         )
         try:
@@ -73,7 +74,7 @@ async def batch_lookup_asns(self, asns, options):
     async def batch_lookup_ips(self, ips, options):
         response = await self._request_with_retry(
             'POST',
-            self._build_base_url('', options),
+            self._build_base_url('', _prefix_batch_fields(options)),
             data=json.dumps(ips)
         )
         try:
@@ -89,7 +90,7 @@ async def batch_lookup_ips(self, ips, options):
     async def batch_parse_user_agents(self, user_agents, options):
         response = await self._request_with_retry(
             'POST',
-            self._build_base_url('user_agent', options),
+            self._build_base_url('user_agent', _prefix_batch_fields(options)),
             data=json.dumps(user_agents)
         )
         try:
@@ -239,6 +240,10 @@ async def batch_request(self, items, request_handler_func, **options):
             )
 
         fresh_item_info = response.data
+        if len(fresh_item_info) != len(cache_misses):
+            raise ClientError(
+                "Batch response contained {} results for {} requested items.".format(
+                    len(fresh_item_info), len(cache_misses)))
         j = 0
         k = 0
 

ipregistry/core.py

@@ -17,7 +17,8 @@
 
 from .cache import IpregistryCache, NoCache
 from .json import AutonomousSystem, IpInfo
-from .model import IpregistryLookupError, ApiResponse, ApiResponseCredits, ApiResponseThrottling
+from .model import (ApiResponse, ApiResponseCredits, ApiResponseThrottling, ClientError,
+                    IpregistryLookupError)
 from .request import DefaultRequestHandler, IpregistryRequestHandler
 
 MAX_BATCH_SIZE = 1024
@@ -124,6 +125,10 @@ def batch_request(self, items, request_handler_func, **options):
             )
 
         fresh_item_info = response.data
+        if len(fresh_item_info) != len(cache_misses):
+            raise ClientError(
+                "Batch response contained {} results for {} requested items.".format(
+                    len(fresh_item_info), len(cache_misses)))
         j = 0
         k = 0
 

ipregistry/request.py

@@ -71,6 +71,25 @@ def _raise_api_error(response):
     raise ApiError(code, message, resolution)
 
 
+def _prefix_batch_fields(options):
+    """Return options with the fields selection adjusted for batch endpoints.
+
+    Batch responses wrap entries in a 'results' array and the API applies the
+    fields filter to the whole response body, so each field must be prefixed
+    with 'results.' to select fields of the individual entries."""
+    fields = options.get('fields')
+    if not fields:
+        return options
+
+    result = dict(options)
+    result['fields'] = ','.join(
+        field if field.startswith('results.') else 'results.' + field
+        for field in (part.strip() for part in str(fields).split(','))
+        if field
+    )
+    return result
+
+
 def _is_retryable_status(config, status_code):
     if status_code == 429:
         return config.retry_on_too_many_requests
@@ -155,7 +174,7 @@ def close(self):
     def batch_lookup_asns(self, asns, options):
         response = self._request_with_retry(
             'POST',
-            self._build_base_url('', options),
+            self._build_base_url('', _prefix_batch_fields(options)),
             data=json.dumps(["AS" + str(asn) for asn in asns])
         )
         try:
@@ -171,7 +190,7 @@ def batch_lookup_asns(self, asns, options):
     def batch_lookup_ips(self, ips, options):
         response = self._request_with_retry(
             'POST',
-            self._build_base_url('', options),
+            self._build_base_url('', _prefix_batch_fields(options)),
             data=json.dumps(ips)
         )
         try:
@@ -187,7 +206,7 @@ def batch_lookup_ips(self, ips, options):
     def batch_parse_user_agents(self, user_agents, options):
         response = self._request_with_retry(
             'POST',
-            self._build_base_url('user_agent', options),
+            self._build_base_url('user_agent', _prefix_batch_fields(options)),
             data=json.dumps(user_agents)
         )
         try:

tests/test_async_client.py

@@ -169,6 +169,35 @@ def responder(request):
             self.assertEqual('1.1.1.1', response.data[0].ip)
             self.assertEqual('INVALID_IP_ADDRESS', response.data[1].code)
 
+    async def test_batch_lookup_fields_are_prefixed(self):
+        """
+        Test that async batch lookups prefix the fields selection with 'results.'
+        """
+        seen_queries = []
+
+        def responder(request):
+            seen_queries.append(request.url.query.decode())
+            return httpx.Response(200, json={'results': [{'ip': '8.8.8.8'}]})
+
+        async with build_client(responder) as client:
+            await client.batch_lookup_ips(['8.8.8.8'], fields='ip,location')
+            self.assertIn('results.ip', seen_queries[0])
+            self.assertIn('results.location', seen_queries[0])
+            self.assertNotIn('fields=ip', seen_queries[0])
+
+    async def test_batch_response_length_mismatch_raises_client_error(self):
+        """
+        Test that a batch response with fewer results than requested items
+        raises a ClientError instead of crashing with an IndexError
+        """
+        def responder(request):
+            return httpx.Response(200, json={})
+
+        async with build_client(responder) as client:
+            with self.assertRaises(ClientError) as context:
+                await client.batch_lookup_ips(['1.1.1.1', '8.8.8.8'])
+            self.assertIn('0 results for 2', str(context.exception))
+
     async def test_missing_httpx_raises_helpful_error(self):
         """
         Test that constructing the handler without httpx raises a helpful error

tests/test_client.py

@@ -114,6 +114,20 @@ def test_batch_lookup_no_split_within_limit(self):
         client.batch_lookup_ips(['1.1.1.1', '1.1.1.2'])
         self.assertEqual(1, handler.calls)
 
+    def test_batch_response_length_mismatch_raises_client_error(self):
+        """
+        Test that a batch response with fewer results than requested items
+        raises a ClientError instead of crashing with an IndexError
+        """
+        class ShortResponseHandler(CountingRequestHandler):
+            def batch_lookup_ips(self, ips, options):
+                return ApiResponse(ApiResponseCredits(), [], ApiResponseThrottling())
+
+        client = IpregistryClient("tryout", requestHandler=ShortResponseHandler())
+        with self.assertRaises(ClientError) as context:
+            client.batch_lookup_ips(['1.1.1.1', '8.8.8.8'])
+        self.assertIn('0 results for 2', str(context.exception))
+
     def test_max_batch_size_capped(self):
         """
         Test that max_batch_size cannot exceed the API limit of 1024
@@ -220,6 +234,16 @@ def test_batch_lookup_ips(self):
         self.assertEqual('INVALID_IP_ADDRESS', response.data[1].code)
         self.assertEqual(True, isinstance(response.data[2], IpInfo))
 
+    def test_batch_lookup_ips_with_fields(self):
+        """
+        Test that batch lookups support field selection
+        """
+        client = IpregistryClient(os.getenv('IPREGISTRY_API_KEY'))
+        response = client.batch_lookup_ips(['8.8.8.8'], fields='ip,location')
+        self.assertEqual('8.8.8.8', response.data[0].ip)
+        self.assertEqual('US', response.data[0].location.country.code)
+        self.assertEqual(None, response.data[0].security)
+
     def test_batch_parse_user_agents(self):
         """
         Test batch parse user agents

tests/test_request.py

@@ -50,6 +50,37 @@ def test_build_base_url_option_value_types(self):
         url = handler._build_base_url('8.8.8.8', {'hostname': True, 'fields': 'location', 'n': 5})
         self.assertEqual('https://api.ipregistry.co/8.8.8.8?hostname=true&fields=location&n=5', url)
 
+    def test_batch_fields_are_prefixed(self):
+        """
+        Test that batch lookups prefix the fields selection with 'results.'
+        as the API applies the filter to the whole batch response body
+        """
+        handler, adapter = build_mocked_handler()
+        adapter.register_uri('POST', 'https://api.ipregistry.co/',
+                             json={'results': [{'ip': '8.8.8.8'}]})
+        handler.batch_lookup_ips(['8.8.8.8'], {'fields': 'ip, location,results.security'})
+        self.assertEqual(['results.ip,results.location,results.security'],
+                         adapter.last_request.qs['fields'])
+
+    def test_batch_parse_user_agents_fields_are_prefixed(self):
+        """
+        Test that batch user-agent parsing prefixes the fields selection too
+        """
+        handler, adapter = build_mocked_handler()
+        adapter.register_uri('POST', 'https://api.ipregistry.co/user_agent',
+                             json={'results': [{'name': 'Chrome'}]})
+        handler.batch_parse_user_agents(['Mozilla/5.0'], {'fields': 'name'})
+        self.assertEqual(['results.name'], adapter.last_request.qs['fields'])
+
+    def test_single_lookup_fields_are_not_prefixed(self):
+        """
+        Test that single lookups keep the fields selection unchanged
+        """
+        handler, adapter = build_mocked_handler()
+        adapter.register_uri('GET', 'https://api.ipregistry.co/8.8.8.8', json={'ip': '8.8.8.8'})
+        handler.lookup_ip('8.8.8.8', {'fields': 'ip,location'})
+        self.assertEqual(['ip,location'], adapter.last_request.qs['fields'])
+
     def test_create_api_error_json_response(self):
         """
         Test that a JSON error response raises an ApiError with its fields