class CheckManagerTest(TestCase):
"""
Base class for all check manager tests.
"""
def setUp(self):
"""
initialize a manager.
"""
self.manager = Linklist()
self.name = 'testlinklist'
self.id = str(uuid4())
self.ids = [self.id]
self.document_content = {
'id': self.id,
'name': self.name,
'linklist': ['http://canopsis.org'],
'mfilter': '{"$and": [{"connector": "collectd"}]}'
}
def clean(self):
self.manager.remove(self.ids)
def get_linklist(self):
return self.manager.find(ids=self.ids)
def linklist_count_equals(self, count):
result = list(self.get_linklist())
if DEBUG:
print result
result = len(result)
self.assertEqual(result, count)
class CTXLinklistRegistry(CTXPropRegistry):
"""In charge of ctx linklist information.
"""
__datatype__ = 'linklist' #: default datatype name
def __init__(self, *args, **kwargs):
super(CTXLinklistRegistry, self).__init__(*args, **kwargs)
self.manager = Linklist()
self.events = MongoStorage(table='events')
self.context = Context()
def _get_documents(self, ids, query):
"""Get documents related to input ids and query.
:param list ids: entity ids. If None, get all documents.
:param dict query: additional selection query.
:return: list of documents.
:rtype: list
"""
result = []
# get entity id field name
ctx_id_field = self._ctx_id_field()
# get a set of entity ids for execution speed reasons
if ids is not None:
ids = set(ids)
# get documents
docs = self.manager.find(_filter=query)
for doc in docs:
try:
mfilter = loads(doc['mfilter'])
except Exception:
pass
else: # get entities from events
events = self.events.find_elements(query=mfilter)
for event in events:
entity = self.context.get_entity(event)
entity_id = self.context.get_entity_id(entity)
if ids is None or entity_id in ids:
doc[ctx_id_field] = entity_id # add eid to the doc
result.append(doc)
return result
def _get(self, ids, query, *args, **kwargs):
return self._get_documents(ids=ids, query=query)
def _delete(self, ids, query, *args, **kwargs):
result = self._get_documents(ids=ids, query=query)
ids = [doc['_id'] for doc in result]
self.manager.remove(ids=ids)
return result
def ids(self, query=None):
result = []
documents = self.manager.find(_filter=query)
for document in documents:
try:
mfilter = loads(document['mfilter'])
except Exception:
pass
else:
# get entities from events
events = self.events.find_elements(query=mfilter)
for event in events:
entity = self.context.get_entity(event)
entity_id = self.context.get_entity_id(entity)
result.append(entity_id)
return result
