def setUp(self):
self.broker = SCIMMABroker()
self.expected_keys = [
'id', 'identifier', 'timestamp', 'topic', 'right_ascension',
'declination', 'right_ascension_sexagesimal',
'declination_sexagesimal', 'parsed_message', 'raw_message',
'created', 'modified'
]
def test_submit_upstream_alert_failure(self):
"""Test that a KafkaException results in an AlertSubmissionException."""
t = Target.objects.create(name='test name', ra=1, dec=2)
with patch('tom_scimma.scimma.Stream.open', mock_open(read_data='data')) as mock_stream:
mock_stream().write.side_effect = KafkaException()
with self.assertRaises(AlertSubmissionException):
SCIMMABroker().submit_upstream_alert(target=t, observation_record=None)
def test_submit_upstream_alert_no_topic(self):
"""Test that submission with no topic succeeds, using the default topic."""
t = Target.objects.create(name='test name', ra=1, dec=2)
with patch('tom_scimma.scimma.Stream.open', mock_open(read_data='data')) as mock_stream:
SCIMMABroker().submit_upstream_alert(target=t, observation_record=None, topic=None)
mock_stream.assert_called_once_with('kafka://test_url:9092/default', 'w')
mock_stream().write.assert_called_with({'type': 'target', 'target_name': t.name, 'ra': t.ra, 'dec': t.dec})
def test_to_generic_alert_lvc_topic(self):
"""Test the to_generic_alert logic for lvc-counterpart alerts. Should result in the inclusion of score."""
test_alert = self.alerts[0]
test_alert['topic'] = 'lvc.lvc-counterpart'
test_alert['message'] = {'rank': '3'}
generic_alert = SCIMMABroker().to_generic_alert(test_alert)
# NOTE: The string is hardcoded as a sanity check to ensure that the string is reviewed if it changes
self.assertEqual(generic_alert.url, f'{SCIMMA_API_URL}/alerts/{test_alert["id"]}')
self.assertEqual(generic_alert.score, '3')
def test_to_generic_alert_any_topic(self):
"""Test the SCIMMA-specific to_generic_alert logic."""
test_alert = self.alerts[0]
test_alert['topic'] = 'gcn'
test_alert['message'] = {'rank': '3'}
generic_alert = SCIMMABroker().to_generic_alert(test_alert)
# NOTE: The string is hardcoded as a sanity check to ensure that the string is reviewed if it changes
self.assertEqual(generic_alert.url, f'{SCIMMA_API_URL}/alerts/{test_alert["id"]}')
self.assertEqual(generic_alert.score, '')
def test_to_target_any_topic(self):
"""Test the SCIMMA-specific to_target logic."""
test_alert = self.alerts[0]
test_alert['topic'] = 'gcn'
SCIMMABroker().to_target(test_alert)
target = Target.objects.filter(name=test_alert['alert_identifier']).first()
self.assertIsInstance(target, Target)
self.assertEqual(target.galactic_lat, None)
self.assertEqual(target.ra, self.alerts[0]['right_ascension'])
self.assertEqual(target.dec, self.alerts[0]['declination'])
def test_fetch_alerts(self, mock_requests_get):
"""Test the SCIMMA-specific fetch_alerts logic."""
mock_response = Response()
mock_response._content = str.encode(json.dumps({'results': self.alerts}))
mock_response.status_code = 200
mock_requests_get.return_value = mock_response
alerts_response = SCIMMABroker().fetch_alerts({})
alerts = [alert for alert in alerts_response]
self.assertEqual(len(alerts), 20)
self.assertEqual(alerts[0], self.alerts[0])
class TestSCIMMAModuleCanary(TestCase):
"""NOTE: To run these tests in your venv: python ./tom_scimma/tests/run_tests.py"""
def setUp(self):
self.broker = SCIMMABroker()
self.expected_keys = [
'id', 'identifier', 'timestamp', 'topic', 'right_ascension',
'declination', 'right_ascension_sexagesimal',
'declination_sexagesimal', 'parsed_message', 'raw_message',
'created', 'modified'
]
def test_boilerplate(self):
self.assertTrue(True)
def test_fetch_alerts(self):
"""Test fetch_alerts"""
response = self.broker.fetch_alerts({'topic': 1})
alerts = []
for alert in response:
alerts.append(alert)
for key in self.expected_keys:
self.assertTrue(key in alert.keys())
self.assertEqual(len(alerts), 20)
def test_fetch_alert(self):
"""Test fetch_alert"""
alert = self.broker.fetch_alert(1000)
self.assertAlmostEqual(alert['right_ascension'], 55.25, 2)
for key in self.expected_keys:
self.assertTrue(key in alert.keys())
def test_submit_upstream_alert(self):
"""Test submit_upstream_alert"""
t = Target.objects.create(name='canary test target', ra=1, dec=2)
response = self.broker.submit_upstream_alert(target=t,
observation_record=None,
topic='TOMToolkit.test')
self.assertTrue(response)
def test_to_target_lvc_topic(self):
"""
Test the to_target logic for lvc-counterpart alerts. Should result in the inclusion of galactic coordinates.
"""
test_alert = self.alerts[0]
test_alert['topic'] = 'lvc-counterpart'
test_alert['message']['gal_coords'] = '336.99,-45.74 [deg] galactic lon,lat of the counterpart'
SCIMMABroker().to_target(test_alert)
target = Target.objects.filter(name=test_alert['alert_identifier']).first()
self.assertIsInstance(target, Target)
self.assertEqual(target.galactic_lat, -45.74)
self.assertEqual(target.ra, self.alerts[0]['right_ascension'])
self.assertEqual(target.dec, self.alerts[0]['declination'])
def test_submit_upstream_alert_no_default_topic(self, mock_stream):
"""Test that submission with no topic and no default topic fails."""
with self.assertRaises(AlertSubmissionException):
SCIMMABroker().submit_upstream_alert(target=None, observation_record=None)