Skip to content

Commit

Permalink
Merge pull request #38 from sam-kleiner/main
Browse files Browse the repository at this point in the history
Save and load json option
  • Loading branch information
meeb authored Aug 1, 2024
2 parents f1a3b74 + f694932 commit 280fd0d
Show file tree
Hide file tree
Showing 2 changed files with 147 additions and 23 deletions.
151 changes: 133 additions & 18 deletions tests/test_public_interface.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import unittest
from datetime import datetime, timezone
from ipaddress import IPv4Network, IPv6Network
Expand Down Expand Up @@ -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')
Expand Down Expand Up @@ -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 ###################
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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())
19 changes: 14 additions & 5 deletions whoisit/bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 280fd0d

Please sign in to comment.