Add helpers for simplifying Atlas usage
whazor opened this issue · 2 comments
whazor commented
In Amundsen I created the following two helpers to work with Atlas Output, I create this issue as we could think about incorporating similar functionality within this library.
A method for transforming an entity collection to a list of entities:
@staticmethod
def _entities(collections: EntityCollection) -> List[Entity]:
"""
Helper method for flattening all collections from {collections}
:return: list of all entities
"""
entities: List[Entity] = []
for collection in collections:
entities.extend(collection.entities)
return entities
Transforming a entity_bulk into a dictionary from 'guids' to entities:
dbs_list = self._entities(self.atlas.entity_bulk(guid=db_ids)) if len(db_ids) > 0 else []
dbs_dict: Dict[str, Entity] = {db.guid: db for db in dbs_list}
whazor commented
We could think of alternative functionality that retrieves the same result as above.
whazor commented
Another method is for mocking responses from Atlas:
def recursive_mock(start: any):
"""
The atlas client allows retrieval of data via __getattr__.
That is why we build this method to recursively mock dictionary's to add
the __getattr__ and to convert them into MagicMock.
:param start: dictionary, list, or other
:return: MagicMock, list with mocked items, or other
"""
if isinstance(start, dict):
dict_mock = MagicMock()
dict_mock.__getitem__.side_effect = start.__getitem__
dict_mock.__iter__.side_effect = start.__iter__
dict_mock.__contains__.side_effect = start.__contains__
dict_mock.get.side_effect = start.get
for key, value in start.items():
value_mock = recursive_mock(value)
dict_mock.__setattr__(key, value_mock)
start[key] = value_mock
return dict_mock
elif isinstance(start, (list,)):
return list(map(recursive_mock, start))
else:
return start
Example usage:
result = recursive_mock({
'entity': {
'name': 'value'
}
})
where you can query with both result.entity.['name']
or result['entity'].name