Added handling of fields as strings

libchickadee/backends/__init__.py

@@ -62,13 +62,18 @@ def query(self, data):
         Args:
             data (list, tuple, set, str): One or more IPs to resolve
 
-        Yield:
-            (dict) request data iterator
+        Returns:
+            (list) List of collected records.
 
         Example:
             >>> records = ['1.1.1.1', '2.2.2.2']
             >>> resolver = ResolverBase()
             >>> resolved_data = resolver.query(records)
+            >>> print(resolved_data)
+            [
+             {"query": "1.1.1.1", "country": "Australia", ...},
+             {"query": "2.2.2.2", "country": "France", ...}
+            ]
         """
 
         self.data = data

libchickadee/backends/ipapi.py

@@ -182,7 +182,7 @@ def batch(self):
 
         for x in orig_recs:
             params = {
-                'fields': ','.join(self.fields),
+                'fields': ','.join(self.fields) if isinstance(self.fields, list) else self.fields,
                 'lang': self.lang,
             }
             if self.api_key:
@@ -237,7 +237,7 @@ def single(self):
         )
         if rdata.status_code == 200:
             self.rate_limit(rdata.headers)
-            return rdata.json()
+            return [rdata.json()]
         elif rdata.status_code == 429:
             self.rate_limit(rdata.headers)
             self.sleeper()

libchickadee/chickadee.py

@@ -194,7 +194,7 @@ def __init__(self, outformat='json', outfile=sys.stdout, fields=_FIELDS):
         self.input_data = None
         self.outformat = outformat
         self.outfile = outfile
-        self.fields = fields
+        self.fields = fields if isinstance(fields, list) else fields.split(',')
         self.force_single = False
         self.ignore_bogon = True
         self.lang = 'en'
@@ -380,7 +380,7 @@ def resolve(self, data_dict, api_key=None):
 
             for element in data:
                 resolver.data = element
-                results.append(resolver.single())
+                results += resolver.single()
         else:
             results = resolver.query(distinct_ips)
 

libchickadee/test/test_backend_ipapi.py

@@ -72,7 +72,7 @@ def test_ipapi_resolve_single(self, mock_query):
             mock_query.return_value = MockResponse(json_data=self.expected_result[count], status_code=200)
             self.resolver.data = ip
             data = self.resolver.single()
-            self.assertEqual(data, self.expected_result[count])
+            self.assertEqual(data, [self.expected_result[count]])
 
     @patch("libchickadee.backends.ipapi.requests.post")
     def test_ipapi_resolve_batch(self, mock_query):
@@ -105,7 +105,7 @@ def test_ipapi_resolve_single_field(self, mock_query):
     def test_ipapi_rate_limiting(self, mock_get, mock_post):
         single = {
             "test_data": self.test_data_ips[1],
-            "expected_data": self.expected_result[1],
+            "expected_data": [self.expected_result[1]],
             "mock_data": [
                 MockResponse(json_data={}, status_code=429, rl='0', ttl='2'),
                 MockResponse(json_data=self.expected_result[1], status_code=200, rl='0', ttl='0')

libchickadee/test/test_chickadee.py

@@ -188,7 +188,7 @@ def test_chickadee_csv_str(self, mock_query):
         chickadee = Chickadee()
         chickadee.ignore_bogon = False
         chickadee.fields = self.fields
-        mock_query.return_value = self.expected_result.copy()
+        mock_query.return_value = self.expected_result
         data = chickadee.run(','.join(self.test_data_ips))
         res = [x for x in data]
         self.assertCountEqual(res, self.expected_result)
@@ -202,7 +202,7 @@ def __init__(self, *args, **kwargs):
                 self.data = None
 
             def single(self):
-                return [x for x in expected_results if x['query'] == self.data][0]
+                return [x for x in expected_results if x['query'] == self.data]
 
         chickadee = Chickadee()
         chickadee.ignore_bogon = False
@@ -222,6 +222,13 @@ def test_improper_type(self):
             failed = True
         self.assertTrue(failed)
 
+    @patch("libchickadee.backends.ipapi.Resolver.batch")
+    def test_manual_run(self, mock_query):
+        chick = Chickadee(fields=self.fields)
+        mock_query.return_value = [self.expected_result[1]]
+        actual = chick.run(self.test_data_ips[1])
+        self.assertDictEqual(self.expected_result[1], actual[0])
+
 
 class ChickadeeFileTestCase(unittest.TestCase):
     """Chickadee script tests."""