def test_chickadee_csv_str(self, mock_query):
"""Batch Query Method Test"""
chickadee = Chickadee()
chickadee.ignore_bogon = False
chickadee.fields = self.fields
mock_query.return_value = self.expected_result
data = chickadee.run(','.join(self.test_data_ips))
self.assertCountEqual(data, self.expected_result)
def test_improper_type(self):
"""Test error handling when the wrong data type is provided via API"""
failed = False
try:
# Provide improper data type
Chickadee.str_handler(['test'])
except TypeError:
failed = True
self.assertTrue(failed)
def test_chickadee_single(self, mock_batch):
"""Query Method Test"""
for count, ip in enumerate(self.test_data_ips):
chickadee = Chickadee()
chickadee.ignore_bogon = False
chickadee.fields = self.fields
mock_batch.return_value = [self.expected_result[count]]
data = chickadee.run(ip)
self.assertEqual(data, [self.expected_result[count]])
def test_ipapi_resolve_query_xlsx_file(self, mock_query):
"""Batch Query Method Test"""
chickadee = Chickadee()
chickadee.ignore_bogon = False
chickadee.fields = self.fields
mock_query.return_value = self.xlsx_data_results
data = chickadee.run(os.path.join(self.test_data_dir, 'test_ips.xlsx'))
batch_result = [] # No reverse field
for item in self.xlsx_data_results:
batch_result.append(item)
self.assertCountEqual(data, batch_result)
def test_chickadee_force_single(self):
"""Batch Query Method Test"""
expected_results = self.expected_result
class MockResolver:
"""Fake resolver"""
def __init__(self, *args, **kwargs):
"""Set defaults"""
self.data = None
def single(self):
"""Mock the single method, returning a list of data."""
return [x for x in expected_results if x['query'] == self.data]
chickadee = Chickadee()
chickadee.ignore_bogon = False
chickadee.force_single = True
chickadee.fields = self.fields
with patch("libchickadee.chickadee.ipapi.Resolver", MockResolver):
data = chickadee.run(','.join(self.test_data_ips))
self.assertCountEqual(data, self.expected_result)
def test_no_resolve(self, resolve='No resolve'):
"""Confirm the handling when no resolution is requested"""
results = [{
'query': '10.0.1.2',
'count': 1,
'message': resolve
}, {
'query': '8.8.8.8',
'count': 1,
'message': resolve
}, {
'query': '2001:4860:4860::8888',
'count': 1,
'message': resolve
}]
for count, ip in enumerate(self.test_data_ips):
chickadee = Chickadee()
chickadee.fields = self.fields
chickadee.resolve_ips = False
chickadee.ignore_bogon = False
data = chickadee.run(ip)
self.assertEqual(data, [results[count]])
def test_file_handler_stream(self):
"""Validate the extraction of IP addresses when input is provided via stdin"""
stream = io.TextIOWrapper(io.StringIO("test 1.1.1.1 ip"))
ips = Chickadee.file_handler(stream, ignore_bogon=True)
self.assertDictEqual(ips, {"1.1.1.1": 1})
def test_ipapi_resolve_query_folder(self, mock_query):
"""Batch Query Method Test"""
expected = [{
"country": "Australia",
"org": "",
"as": "AS13335 Cloudflare, Inc.",
"proxy": False,
"query": "1.1.1.1",
"count": 6
}, {
"query": "10.0.1.2",
"count": 3
}, {
"country": "United States",
"org": "Level 3",
"as": "AS15169 Google LLC",
"proxy": False,
"query": "8.8.8.8",
"count": 3
}, {
"country": "United States",
"org": "Google LLC",
"as": "AS15169 Google LLC",
"proxy": False,
"query": "2001:4860:4860::8888",
"count": 3
}, {
"country": "United States",
"org": "Informs",
"as": "AS3356 Level 3 Parent, LLC",
"proxy": True,
"query": "4.4.4.4",
"count": 3
}, {
"country": "United States",
"org": "Google LLC",
"as": "AS15169 Google LLC",
"proxy": False,
"query": "2001:4860:4860::8844",
"count": 4
}, {
"country": "United States",
"org": "Google LLC",
"as": "AS15169 Google LLC",
"proxy": False,
"query": "2001:4860:4860::8888",
"count": 3
}, {
"country": "France",
"org": "",
"as": "AS3215 Orange S.A.",
"proxy": True,
"query": "2.2.2.2",
"count": 3
}, {
"country": "United States",
"org": "Google LLC",
"as": "AS15169 Google LLC",
"proxy": False,
"query": "2001:4860:4860::8844",
"count": 4
}]
mock_query.return_value = expected
chickadee = Chickadee()
chickadee.ignore_bogon = False
chickadee.fields = self.fields
data = chickadee.run(self.test_data_dir)
self.assertCountEqual(data, expected)
def test_manual_run(self, mock_query):
"""Test the execution of chickadee via module calls"""
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])