diff --git a/tests/test_public_interface.py b/tests/test_public_interface.py index 21a7630..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 @@ -93,15 +94,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') @@ -224,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 ################### @@ -327,15 +385,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") @@ -457,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()) diff --git a/whoisit/bootstrap.py b/whoisit/bootstrap.py index f87437f..5e43ace 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'dict, 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