def handle(self, *args, **options):
brokers = alerts.get_service_classes()
broker_classes = {}
for broker in brokers:
broker_classes[broker] = alerts.get_service_class(broker)()
target = None
if options['target_id']:
try:
targets = [Target.objects.get(pk=options['target_id'])]
except ObjectDoesNotExist:
raise Exception('Invalid target id provided')
else:
sources = ReducedDatumSource.objects.filter(
name__in=broker_classes.keys()).distinct()
targets = Target.objects.filter(id__in=ReducedDatum.objects.filter(
source__in=sources).values_list('target').distinct())
failed_records = {}
for target in targets:
for class_name, clazz in broker_classes.items():
try:
clazz.process_reduced_data(target)
except HTTPError:
failed_records[class_name] = target.id
if len(failed_records) == 0:
return 'Update completed successfully'
else:
return 'Update completed with errors: {0}'.format(
str(failed_records))
class BrokerQueryFilter(django_filters.FilterSet):
broker = django_filters.ChoiceFilter(
choices=[(k, k) for k in get_service_classes().keys()])
name = django_filters.CharFilter(lookup_expr='icontains')
class Meta:
model = BrokerQuery
fields = ['broker', 'name']
def test_get_service_classes_import_error(self):
with self.subTest('Invalid import returns an import error.'):
with patch(
'tom_alerts.alerts.import_module') as mock_import_module:
mock_import_module.side_effect = ImportError()
with self.assertRaisesRegex(
ImportError,
'Could not import tom_alerts.fake_broker.'):
get_service_classes()
with self.subTest('Invalid import returns an attribute error.'):
with patch(
'tom_alerts.alerts.import_module') as mock_import_module:
mock_import_module.side_effect = AttributeError()
with self.assertRaisesRegex(
ImportError,
'Could not import tom_alerts.fake_broker.'):
get_service_classes()
def get_context_data(self, *args, **kwargs):
"""
Adds the brokers available to the TOM to the context dictionary.
:returns: context
:rtype: dict
"""
context = super().get_context_data(*args, **kwargs)
context['installed_brokers'] = get_service_classes()
return context
class BrokerQueryFilter(FilterSet):
"""
Defines the available fields for filtering the list of broker queries.
"""
broker = ChoiceFilter(choices=[(k, k)
for k in get_service_classes().keys()])
name = CharFilter(lookup_expr='icontains')
class Meta:
model = BrokerQuery
fields = ['broker', 'name']
def get_context_data(self, *args, **kwargs):
"""
Adds the brokers available to the TOM to the context dictionary.
:returns: context
:rtype: dict
"""
context = super().get_context_data(*args, **kwargs)
private_brokers = settings.PRIVATE_BROKERS
context['private_brokers'] = private_brokers
context['amon_brokers'] = settings.AMON_BROKERS
installed_brokers = get_service_classes()
if not self.request.user.groups.filter(name='AMON members').exists():
for broker in private_brokers:
if broker in installed_brokers:
del installed_brokers[broker]
context['installed_brokers'] = installed_brokers
return context
def get_context_data(self, *args, **kwargs):
context = super().get_context_data(*args, **kwargs)
context['installed_brokers'] = get_service_classes()
return context