class TestIntegration(unittest.TestCase):
def setUp(self):
data_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), 'data'))
schema_path = os.path.abspath(
os.path.join(os.path.dirname(__file__),
'../schemas/distribution_schema_0p2.avsc'))
r = AlertReader(data_path)
alerts = r.to_list()
conf = load_credentials(tmp=True)
kafka_servers = conf['servers']
p = confluent_kafka.Producer({'bootstrap.servers': kafka_servers})
for alert in alerts:
avro_data = encode_into_avro(alert, schema_path)
topic = get_legal_topic_name(alert['cdsxmatch'])
p.produce(topic, avro_data)
p.flush()
# instantiate an AlertConsumer
mytopics = conf["mytopics"]
myconfig = {
'bootstrap.servers': kafka_servers,
'group_id': conf['group_id']
}
self.consumer = AlertConsumer(mytopics,
myconfig,
schema_path=schema_path)
def test_poll(self):
topic, alert, key = self.consumer.poll()
self.assertIsNotNone(alert)
self.assertTrue(fastavro.validate(alert, self.consumer._parsed_schema))
def test_consume(self):
num_messages = 1
alerts = self.consumer.consume(num_messages)
self.assertEqual(len(alerts), num_messages)
def test_topics(self):
topics = self.consumer.available_topics()
self.assertTrue('rrlyr' in topics.keys())
def test_broker_name(self):
brokers = self.consumer.available_brokers()
self.assertTrue(0 in brokers.keys())
def tearDown(self):
self.consumer.close()
def main():
""" """
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
'--display',
action='store_true',
help="If specified, print on screen information about incoming alert.")
parser.add_argument(
'-limit',
type=int,
default=None,
help="If specified, download only `limit` alerts. Default is None.")
parser.add_argument(
'--available_topics',
action='store_true',
help="If specified, print on screen information about available topics."
)
parser.add_argument(
'--save',
action='store_true',
help="If specified, save alert data on disk (Avro). See also -outdir.")
parser.add_argument(
'-outdir',
type=str,
default='.',
help="Folder to store incoming alerts if --save is set. It must exist."
)
parser.add_argument(
'-schema',
type=str,
default=None,
help=
"Avro schema to decode the incoming alerts. Default is None (latest version downloaded from server)"
)
args = parser.parse_args(None)
# load user configuration
conf = load_credentials()
myconfig = {
"username": conf['username'],
'bootstrap.servers': conf['servers'],
'group_id': conf['group_id']
}
if conf['password'] is not None:
myconfig['password'] = conf['password']
# Instantiate a consumer
if args.schema is None:
schema = None
else:
schema = args.schema
consumer = AlertConsumer(conf['mytopics'], myconfig, schema_path=schema)
if args.available_topics:
print(consumer.available_topics().keys())
sys.exit(0)
# Time to wait before polling again if no alerts
maxtimeout = conf['maxtimeout']
# infinite loop
maxpoll = args.limit if args.limit else 1e10
try:
poll_number = 0
while poll_number < maxpoll:
if args.save:
# Save alerts on disk
topic, alert, key = consumer.poll_and_write(outdir=args.outdir,
timeout=maxtimeout,
overwrite=True)
else:
# TODO: this is useless to get it and done nothing
# why not thinking about handler like Comet?
topic, alert, key = consumer.poll(timeout=maxtimeout)
if topic is not None:
poll_number += 1
if args.display and topic is not None:
utc = time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime())
table = [
[
alert['timestamp'], utc, topic, alert['objectId'],
alert['cdsxmatch'], alert['rfscore']
],
]
headers = [
'Emitted at (UTC)', 'Received at (UTC)', 'Topic',
'objectId', 'Simbad', 'RF score'
]
print(tabulate(table, headers, tablefmt="pretty"))
elif args.display:
print('No alerts the last {} seconds'.format(maxtimeout))
except KeyboardInterrupt:
sys.stderr.write('%% Aborted by user\n')
finally:
consumer.close()