jpoullet2000/atlasclient

Add helpers for simplifying Atlas usage

whazor opened this issue · 2 comments

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}

We could think of alternative functionality that retrieves the same result as above.

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