From d2345fa26381c92f0d4acc7ce3c22edd85983f58 Mon Sep 17 00:00:00 2001 From: Sam Kleiner Date: Thu, 1 Aug 2024 12:49:05 -0400 Subject: [PATCH 1/4] add save and load as json option --- whoisit/bootstrap.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/whoisit/bootstrap.py b/whoisit/bootstrap.py index f87437f..72e87d3 100644 --- a/whoisit/bootstrap.py +++ b/whoisit/bootstrap.py @@ -137,24 +137,33 @@ def _bootstrap(self, overrides=False, allow_insecure=False): raise BootstrapError(f'Failed to load some bootstrap data, ' f'missing data: {items_missing}') - def save_bootstrap_data(self): + def save_bootstrap_data(self, as_json=True): if not self.is_bootstrapped(): raise BootstrapError('No bootstrap data is loaded') rtn = {'timestamp': self._bootstrap_timestamp} for name, data in self._data.items(): rtn[name] = data - return json.dumps(rtn) - def load_bootstrap_data(self, data, overrides=False, allow_insecure=False): + if as_json: + return json.dumps(rtn) + + return rtn + + + def load_bootstrap_data(self, data, overrides=False, allow_insecure=False, from_json=True): if self.is_bootstrapped(): raise BootstrapError('Already bootstrapped, cannot load more data') - if not isinstance(data, str): + if not isinstance(data, str) and from_json: + raise BootstrapError(f'Unable to load bootstrap data, data must be a ' + f'string, got: {type(data)}') + elif not isinstance(data, dict) and not from_json: raise BootstrapError(f'Unable to load bootstrap data, data must be a ' f'string, got: {type(data)}') self._use_iana_overrides = bool(overrides) self._allow_insecure = bool(allow_insecure) try: - data = json.loads(data) + if from_json: + data = json.loads(data) except Exception as e: raise BootstrapError(f'Unable to load bootstrap data, failed to parse ' f'as JSON: {e}') from e From 5f6478568faf9197d7ab07df60dc70d8253057bc Mon Sep 17 00:00:00 2001 From: Sam Kleiner Date: Thu, 1 Aug 2024 12:50:23 -0400 Subject: [PATCH 2/4] format --- tests/test_public_interface.py | 42 +++++++++++++++++++--------------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/tests/test_public_interface.py b/tests/test_public_interface.py index 21a7630..13f6bf4 100644 --- a/tests/test_public_interface.py +++ b/tests/test_public_interface.py @@ -93,15 +93,18 @@ def test_domain_interface_1(self): self.assertTrue(len(resp["entities"]["registrar"]) > 0) self.assertTrue(isinstance(resp["entities"]["registrant"], list)) self.assertTrue(len(resp["entities"]["registrant"]) > 0) - self.assertEqual(resp["entities"]["registrant"][0]['address'], { - 'country': "US", - 'ext_address': '', - 'locality': '', - 'po_box': '', - 'postal_code': '', - 'region': 'CA', - 'street_address': '', - }) + self.assertEqual( + resp["entities"]["registrant"][0]["address"], + { + "country": "US", + "ext_address": "", + "locality": "", + "po_box": "", + "postal_code": "", + "region": "CA", + "street_address": "", + }, + ) @responses.activate # @_recorder.record(file_path=RESPONSES / 'ip-v4-1.yaml') @@ -327,15 +330,18 @@ async def test_domain_interface_1(self): self.assertTrue(len(resp["entities"]["registrar"]) > 0) self.assertTrue(isinstance(resp["entities"]["registrant"], list)) self.assertTrue(len(resp["entities"]["registrant"]) > 0) - self.assertEqual(resp["entities"]["registrant"][0]['address'], { - 'country': "US", - 'ext_address': '', - 'locality': '', - 'po_box': '', - 'postal_code': '', - 'region': 'CA', - 'street_address': '', - }) + self.assertEqual( + resp["entities"]["registrant"][0]["address"], + { + "country": "US", + "ext_address": "", + "locality": "", + "po_box": "", + "postal_code": "", + "region": "CA", + "street_address": "", + }, + ) @pytest.mark.asyncio @pytest.mark.usefixtures("mock_httpx") From 8b66a1547a89f61e06c85f861bf7800bdb7a568a Mon Sep 17 00:00:00 2001 From: Sam Kleiner Date: Thu, 1 Aug 2024 12:50:38 -0400 Subject: [PATCH 3/4] add save and load tests --- tests/test_public_interface.py | 109 +++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) diff --git a/tests/test_public_interface.py b/tests/test_public_interface.py index 13f6bf4..09d1be6 100644 --- a/tests/test_public_interface.py +++ b/tests/test_public_interface.py @@ -1,3 +1,4 @@ +import json import unittest from datetime import datetime, timezone from ipaddress import IPv4Network, IPv6Network @@ -227,6 +228,60 @@ def test_entity_interface_1(self): self.assertTrue(isinstance(resp["entities"]["technical"], list)) self.assertTrue(len(resp["entities"]["technical"]) > 0) + @responses.activate + # @_recorder.record(file_path=RESPONSES / 'boostrap.yaml') + def test_save_bootstrap_data(self): + responses._add_from_file(RESPONSES / "boostrap.yaml") + whoisit.bootstrap() + self.assertTrue(whoisit.is_bootstrapped()) + + result = whoisit.save_bootstrap_data() + self.assertTrue(isinstance(result, str)) + + data = json.loads(result) + self.assertTrue(bool(data)) + + @responses.activate + # @_recorder.record(file_path=RESPONSES / 'boostrap.yaml') + def test_save_bootstrap_data_no_json(self): + responses._add_from_file(RESPONSES / "boostrap.yaml") + whoisit.bootstrap() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data(as_json=False) + self.assertTrue(isinstance(data, dict)) + self.assertTrue(bool(data)) + + @responses.activate + # @_recorder.record(file_path=RESPONSES / 'boostrap.yaml') + def test_load_bootstrap_data(self): + responses._add_from_file(RESPONSES / "boostrap.yaml") + whoisit.bootstrap() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data() + self.assertTrue(isinstance(data, str)) + + whoisit.clear_bootstrapping() + + whoisit.load_bootstrap_data(data) + self.assertTrue(whoisit.is_bootstrapped()) + + @responses.activate + # @_recorder.record(file_path=RESPONSES / 'boostrap.yaml') + def test_load_bootstrap_data_no_json(self): + responses._add_from_file(RESPONSES / "boostrap.yaml") + whoisit.bootstrap() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data(as_json=False) + self.assertTrue(isinstance(data, dict)) + + whoisit.clear_bootstrapping() + + whoisit.load_bootstrap_data(data, from_json=False) + self.assertTrue(whoisit.is_bootstrapped()) + ############################################# ################### Async ################### @@ -463,3 +518,57 @@ async def test_entity_interface_1(self): ) self.assertTrue(isinstance(resp["entities"]["technical"], list)) self.assertTrue(len(resp["entities"]["technical"]) > 0) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("mock_httpx") + async def test_save_bootstrap_data(self): + load_sync_responses_to_httpx_mock(RESPONSES / "boostrap.yaml", self.httpx_mock) + await whoisit.bootstrap_async() + self.assertTrue(whoisit.is_bootstrapped()) + + result = whoisit.save_bootstrap_data() + self.assertTrue(isinstance(result, str)) + + data = json.loads(result) + self.assertTrue(bool(data)) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("mock_httpx") + async def test_save_bootstrap_data_no_json(self): + load_sync_responses_to_httpx_mock(RESPONSES / "boostrap.yaml", self.httpx_mock) + await whoisit.bootstrap_async() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data(as_json=False) + self.assertTrue(isinstance(data, dict)) + self.assertTrue(bool(data)) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("mock_httpx") + async def test_load_bootstrap_data(self): + load_sync_responses_to_httpx_mock(RESPONSES / "boostrap.yaml", self.httpx_mock) + await whoisit.bootstrap_async() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data() + self.assertTrue(isinstance(data, str)) + + whoisit.clear_bootstrapping() + + whoisit.load_bootstrap_data(data) + self.assertTrue(whoisit.is_bootstrapped()) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("mock_httpx") + async def test_load_bootstrap_data_no_json(self): + load_sync_responses_to_httpx_mock(RESPONSES / "boostrap.yaml", self.httpx_mock) + await whoisit.bootstrap_async() + self.assertTrue(whoisit.is_bootstrapped()) + + data = whoisit.save_bootstrap_data(as_json=False) + self.assertTrue(isinstance(data, dict)) + + whoisit.clear_bootstrapping() + + whoisit.load_bootstrap_data(data, from_json=False) + self.assertTrue(whoisit.is_bootstrapped()) From f6949323559fde85d76ed797e0406e90b86f4dbe Mon Sep 17 00:00:00 2001 From: Sam Kleiner Date: Thu, 1 Aug 2024 12:54:55 -0400 Subject: [PATCH 4/4] fix exception message --- whoisit/bootstrap.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/whoisit/bootstrap.py b/whoisit/bootstrap.py index 72e87d3..5e43ace 100644 --- a/whoisit/bootstrap.py +++ b/whoisit/bootstrap.py @@ -158,7 +158,7 @@ def load_bootstrap_data(self, data, overrides=False, allow_insecure=False, from_ f'string, got: {type(data)}') elif not isinstance(data, dict) and not from_json: raise BootstrapError(f'Unable to load bootstrap data, data must be a ' - f'string, got: {type(data)}') + f'dict, got: {type(data)}') self._use_iana_overrides = bool(overrides) self._allow_insecure = bool(allow_insecure) try: