def test_get_data_with_limit(self):
keys = ('id', ('users', 0, 'emailAddress'))
expected_model_data = ({'id': 1, 'emailAddress': '*****@*****.**'}, )
with mock.patch.object(FakeClient,
'find_fake_table_iter',
return_value=iter(self.model_data)):
mt = ModelTrawler('fake_table', FakeClient())
assert tuple(mt.get_data(keys, limit=1)) == expected_model_data
def test_filter_keys(self):
keys = ('id', ('supplier', 'name'), ('users', 0, 'emailAddress'))
initial_model_dict = self.model_data[0].copy()
expected_model_dict = {
'id': 1,
'name': 'Super cloud supplier',
'emailAddress': '*****@*****.**'
}
mt = ModelTrawler('fake_table', FakeClient())
assert mt._filter_keys(keys)(initial_model_dict) == expected_model_dict
def test_get_data_with_kwargs(self):
_kwargs = {'something_true': True, 'something_false': False}
with mock.patch.object(FakeClient,
'find_fake_table_iter',
return_value=(m for m in ({
'id': 1
}, {
'id': 2
}))) as mock_method:
mt = ModelTrawler('fake_table', FakeClient())
list(mt.get_data(('id', ), limit=1, **_kwargs))
# note that `limit` isn't included as a keyword argument
mock_method.assert_called_once_with(something_true=True,
something_false=False)
def base_model(base_model,
keys,
get_data_kwargs,
client,
logger=None,
limit=None):
"""Fetch all the data for a given Digital Marketplace model from the api.
:param base_model: A Digital Marketplace model (client must have a 'find_{model}_iter' method)
:param keys: The attributes we require of that model.
:param get_data_kwargs: Additional kwargs for the get request the client will make.
:param client: Instantiated Digital Marketplace APIClient
:param logger:
:param limit: Maximum number of requests for the client to perform.
:return: A pandas DataFrame of the requested data. Columns as model attributes, rows as instances.
"""
mt = ModelTrawler(base_model, client)
data = list(mt.get_data(keys=keys, limit=limit, **get_data_kwargs))
if logger:
logger.info('{} {} returned after {}s'.format(len(data), base_model,
mt.get_time_running()))
return pandas.DataFrame(data) if data else pandas.DataFrame(columns=keys)
def test_get_allowed_models_returns_allowed_models(self):
mt = ModelTrawler('fake_table', FakeClient())
assert set(mt._get_allowed_models()) == set(['fake_table'])
def test_cannot_initiate_with_incorrect_model_name(self):
with pytest.raises(AttributeError):
ModelTrawler('real_table', FakeClient())
def test_correct_initiation(self):
mt = ModelTrawler('fake_table', FakeClient())
assert mt.model == 'fake_table'
assert mt.model_iter_method == 'find_fake_table_iter'