def test_subscribe_to_temporary_queue(mockstomp):
"""Test subscribing to a topic (publish-subscribe) and callback functions."""
mock_cb = mock.Mock()
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
known_subscriptions = set()
known_queues = set()
def assert_not_seen_before(ts: TemporarySubscription):
assert ts.subscription_id, "Temporary subscription is missing an ID"
assert (ts.subscription_id
not in known_subscriptions), "Duplicate subscription ID"
assert ts.queue_name, "Temporary queue does not have a name"
assert ts.queue_name not in known_queues, "Duplicate temporary queue name"
known_subscriptions.add(ts.subscription_id)
known_queues.add(ts.queue_name)
print(f"Temporary subscription: {ts}")
mockconn.set_listener.assert_called_once()
listener = mockconn.set_listener.call_args[0][1]
assert listener is not None
ts = {}
for n, queue_hint in enumerate(
("", "", "hint", "hint", "transient.hint", "transient.hint")):
ts[n] = stomp.subscribe_temporary(
channel_hint=queue_hint,
callback=mock_cb,
)
assert_not_seen_before(ts[n])
assert ts[n].queue_name.startswith("transient.")
return
class Master(object):
def __init__(self):
self.transport = StompTransport()
self.transport.connect()
self.heartbeat = HeartBeat(self.transport)
self.heartbeat.start()
self.sender = Sender(self.transport)
self.receiver = Receiver(self.transport)
self.num_messages = 0
def send(self, message):
self.num_messages += 1
self.sender.send(message)
def receive(self):
try:
from time import sleep
while len(self.receiver.message_list) != self.num_messages:
print(len(self.receiver.message_list))
sleep(1)
self.heartbeat.stop()
except KeyboardInterrupt:
self.heartbeat.stop()
raise
return self.receiver.message_list
def test_messages_are_not_serialized_for_raw_transport(mockstomp):
"""Test the raw sending methods."""
banana = '{"entry": [0, "banana"]}'
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp.raw_send(str(mock.sentinel.channel1), banana)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/queue/" + str(mock.sentinel.channel1), banana)
mockconn.send.reset_mock()
stomp.raw_broadcast(str(mock.sentinel.channel2), banana)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel2), banana)
mockconn.send.reset_mock()
stomp.raw_send(str(mock.sentinel.channel), mock.sentinel.unserializable)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == (
"/queue/" + str(mock.sentinel.channel),
mock.sentinel.unserializable,
)
def test_send_broadcast(mockstomp):
"""Test the broadcast sending function."""
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp._broadcast(str(mock.sentinel.channel), mock.sentinel.message)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs.get("headers") in (None, {})
stomp._broadcast(str(mock.sentinel.channel),
mock.sentinel.message,
headers=mock.sentinel.headers)
assert mockconn.send.call_count == 2
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs == {"headers": mock.sentinel.headers}
stomp._broadcast(str(mock.sentinel.channel),
mock.sentinel.message,
delay=123)
assert mockconn.send.call_count == 3
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs["headers"].get("AMQ_SCHEDULED_DELAY") == 123000
def __init__(self, wait_time=5):
self.wait_time = wait_time
self.heartbeat_timestamp = None
self.transport = StompTransport()
self.transport.connect()
self.transport.subscribe_broadcast("heartbeat", self.read_heartbeat)
self.transport.subscribe("outbound", self.read)
self.wait()
def __init__(self):
self.transport = StompTransport()
self.transport.connect()
self.heartbeat = HeartBeat(self.transport)
self.heartbeat.start()
self.sender = Sender(self.transport)
self.receiver = Receiver(self.transport)
self.num_messages = 0
def test_error_handling_when_connecting_to_broker(mockstomp):
"""Test the Stomp connection routine."""
stomp = StompTransport()
mockconn = mockstomp.Connection.return_value
mockconn.connect.side_effect = stomppy.exception.ConnectFailedException()
mockstomp.exception.ConnectFailedException = (
stomppy.exception.ConnectFailedException)
with pytest.raises(workflows.Disconnected):
stomp.connect()
assert not stomp.is_connected()
def __init__(self):
# initialize logging
self.setup_logging()
# change settings when in live mode
default_configuration = "/dls_sw/apps/zocalo/secrets/credentials-testing.cfg"
if "--live" in sys.argv:
self.use_live_infrastructure = True
default_configuration = "/dls_sw/apps/zocalo/secrets/credentials-live.cfg"
else:
self.use_live_infrastructure = False
if os.path.exists(default_configuration):
StompTransport.load_configuration_file(default_configuration)
def test_send_message(mockstomp):
"""Test the message sending function."""
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp._send(str(mock.sentinel.channel), mock.sentinel.message)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/queue/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs.get("headers") == {"persistent": "true"}
stomp._send(
str(mock.sentinel.channel),
mock.sentinel.message,
headers={"hdr": mock.sentinel.header},
delay=123,
)
assert mockconn.send.call_count == 2
args, kwargs = mockconn.send.call_args
assert args == ("/queue/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs == {
"headers": {
"hdr": mock.sentinel.header,
"persistent": "true",
"AMQ_SCHEDULED_DELAY": 123000,
}
}
def test_check_config_file_behaviour(mockstomp):
"""Check that a specified configuration file is read, that command line
parameters have precedence and are passed on to the stomp layer."""
mockconn = mock.Mock()
mockstomp.Connection.return_value = mockconn
parser = optparse.OptionParser()
stomp = StompTransport()
stomp.add_command_line_options(parser)
# Temporarily create an example stomp configuration file
cfgfile = tempfile.NamedTemporaryFile(delete=False)
try:
cfgfile.write("""
# An example stomp configuration file
# Only lines in the [stomp] block will be interpreted
[stomp]
#host = 127.0.0.1
port = 1234
username = someuser
password = somesecret
prefix = namespace
""".encode("utf-8"))
cfgfile.close()
parser.parse_args(
["--stomp-conf", cfgfile.name, "--stomp-user", mock.sentinel.user])
# Command line parameters are shared for all instances
stomp = StompTransport()
stomp.connect()
# Reset configuration for subsequent tests by reloading StompTransport
importlib.reload(workflows.transport.stomp_transport)
globals(
)["StompTransport"] = workflows.transport.stomp_transport.StompTransport
mockstomp.Connection.assert_called_once_with([("localhost", 1234)])
mockconn.connect.assert_called_once_with(mock.sentinel.user,
"somesecret",
wait=False)
assert stomp.get_namespace() == "namespace"
finally:
os.remove(cfgfile.name)
# Loading a non-existing configuration file
with pytest.raises(workflows.Error):
parser.parse_args(["--stomp-conf", ""])
def test_messages_are_serialized_for_transport(mockstomp):
"""Test the message serialization."""
banana = {"entry": [1, 2.0, decimal.Decimal(3), "banana"]}
banana_str = '{"entry": [1, 2.0, 3.0, "banana"]}'
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp.send(str(mock.sentinel.channel1), banana)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/queue/" + str(mock.sentinel.channel1), banana_str)
stomp.broadcast(str(mock.sentinel.channel2), banana)
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel2), banana_str)
with pytest.raises(Exception):
stomp.send(str(mock.sentinel.channel), mock.sentinel.unserializable)
def test_adding_arguments_to_argparser():
"""Check that command line parameters can be added to the parser."""
parser = argparse.ArgumentParser()
StompTransport().add_command_line_options(parser)
result = parser.parse_args([])
assert result.stomp_host
assert result.stomp_port
assert result.stomp_user
assert result.stomp_pass
def test_add_command_line_help_optparse():
"""Check that command line parameters are registered in the parser."""
parser = mock.MagicMock()
StompTransport().add_command_line_options(parser)
parser.add_argument.assert_not_called()
parser.add_option.assert_called()
assert parser.add_option.call_count > 4
for call in parser.add_option.call_args_list:
assert call[1]["action"] == "callback"
def test_nack_message(mockstomp):
"""Test that the _nack function is properly forwarded to stomp."""
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
subid = stomp._subscribe(1,
str(mock.sentinel.channel3),
None,
acknowledgement=True)
stomp._nack(mock.sentinel.messageid, subid)
mockconn.nack.assert_called_once_with(mock.sentinel.messageid, subid)
stomp._nack(mock.sentinel.messageid, subid, transaction=mock.sentinel.txn)
mockconn.nack.assert_called_with(mock.sentinel.messageid,
subid,
transaction=mock.sentinel.txn)
def send_to_stomp_or_defer(message, headers=None):
if not headers:
headers = generate_headers()
if options.verbose:
pprint(message)
if allow_stomp_fallback and options.dropfile:
return write_message_to_dropfile(message, headers)
try:
stomp = StompTransport()
if options.dryrun:
print("Not sending message (running with --dry-run)")
return
stomp.connect()
stomp.send("processing_recipe", message, headers=headers)
except (
KeyboardInterrupt,
SyntaxError,
AssertionError,
AttributeError,
ImportError,
TypeError,
ValueError,
):
raise
except Exception:
if not allow_stomp_fallback:
raise
print("\n\n")
import traceback
traceback.print_exc()
print("\n\nAttempting to store message in fallback location")
write_message_to_dropfile(message, headers)
def test_error_handling_on_broadcast(mockstomp):
"""Unrecoverable errors during broadcasting should mark the connection as disconnected."""
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
mockconn.send.side_effect = stomppy.exception.NotConnectedException()
mockstomp.exception = stomppy.exception
with pytest.raises(workflows.Disconnected):
stomp._broadcast(str(mock.sentinel.channel), mock.sentinel.message)
assert not stomp.is_connected()
def __init__(self):
# initialize logging
self.setup_logging()
self.log.debug('Loading dlstbx workflows plugins')
dlstbx = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
self.log.debug('Loading dlstbx credentials')
# change settings when in live mode
default_configuration = '/dls_sw/apps/zocalo/secrets/credentials-testing.cfg'
if '--live' in sys.argv:
self.use_live_infrastructure = True
default_configuration = '/dls_sw/apps/zocalo/secrets/credentials-live.cfg'
# override default stomp host
from workflows.transport.stomp_transport import StompTransport
try:
StompTransport.load_configuration_file(default_configuration)
except workflows.Error as e:
self.log.warning(e)
class Worker(object):
def __init__(self, wait_time=5):
self.wait_time = wait_time
self.heartbeat_timestamp = None
self.transport = StompTransport()
self.transport.connect()
self.transport.subscribe_broadcast("heartbeat", self.read_heartbeat)
self.transport.subscribe("outbound", self.read)
self.wait()
def timeout(self):
from time import time
if self.heartbeat_timestamp is None:
self.heartbeat_timestamp = time()
else:
if time() - self.heartbeat_timestamp > self.wait_time:
return True
return False
def wait(self):
from time import time, sleep
try:
while not self.timeout():
sleep(1)
except KeyboardInterrupt:
pass
def send(self, message):
print("Sending")
self.transport.send("inbound", message)
def read_heartbeat(self, header, message):
from time import time
self.heartbeat_timestamp = time()
def read(self, header, message):
message = {"result": "Hey Hey!"}
self.send(message)
def test_anonymous_connection(mockstomp):
"""Check that a specified configuration file is read, that command line
parameters have precedence and are passed on to the stomp layer."""
mockconn = mock.Mock()
mockstomp.Connection.return_value = mockconn
parser = optparse.OptionParser()
stomp = StompTransport()
stomp.add_command_line_options(parser)
parser.parse_args(["--stomp-user="******"--stomp-pass="******"StompTransport"] = workflows.transport.stomp_transport.StompTransport
mockconn.connect.assert_called_once_with(wait=False)
def test_broadcast_status(mockstomp, mocktime):
"""Test the status broadcast function."""
mocktime.time.return_value = 20000
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp.broadcast_status({"status": str(mock.sentinel.status)})
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
# expiration should be 15 seconds in the future
assert int(kwargs["headers"]["expires"]) == 1000 * (20000 + 15)
destination, message = args
assert destination.startswith("/topic/transient.status")
statusdict = json.loads(message)
assert statusdict["status"] == str(mock.sentinel.status)
def test_broadcasting_message_with_expiration(time, mockstomp):
"""Test sending a message that expires some time in the future."""
system_time = 1234567.1234567
message_lifetime = 120
expiration_time = int((system_time + message_lifetime) * 1000)
time.time.return_value = system_time
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp._broadcast(str(mock.sentinel.channel),
mock.sentinel.message,
expiration=120)
mockconn.send.assert_called_once()
args, kwargs = mockconn.send.call_args
assert args == ("/topic/" + str(mock.sentinel.channel),
mock.sentinel.message)
assert kwargs.get("headers") == {"expires": expiration_time}
def lazy_pprint(*args, **kwargs):
from pprint import pprint
# pprint(*args, **kwargs)
# def get_output_file(file_path):
# print(os.path.dirname(get_output_file()))
# return os.path.dirname(get_output_file())
if __name__ == '__main__':
default_configuration = '/dls_sw/apps/zocalo/secrets/credentials-live.cfg'
# override default stomp host
try:
StompTransport.load_configuration_file(default_configuration)
except workflows.Error as e:
print("Error: %s\n" % str(e))
# StompTransport.add_command_line_options(parser)
# (options, args) = parser.parse_args(sys.argv[1:])
stomp = StompTransport()
output_file = get_output_file(sys.argv[-1])
message = {
'recipes': [],
'parameters': {},
}
# Build a custom recipe
recipe = {}
for p in processes:
p.join()
class Data(object):
def __init__(self, a, b, c):
self.a = a
self.b = b
self.c = c
if __name__ == "__main__":
import cPickle as pickle
configuration = "/dls_sw/apps/zocalo/secrets/credentials-james.cfg"
StompTransport.load_configuration_file(configuration)
data = Data(1, 2, 3)
message = {"hello": "world", "data": pickle.dumps(data)}
master = Master()
master.send(message)
master.send(message)
master.send(message)
master.send(message)
print(master.receive())
print("\nSubmitted.")
secure=False,
)
def ensure_url(url: str) -> str:
"""Make sure a string URL has a schema, for urllib.parse.urlparse consumption"""
if "://" not in url:
return f"minio://{args.host}"
return url
logger = logging.getLogger()
parser = ArgumentParser(description="Submit an S3 bucket for PIA")
StompTransport.load_configuration_file(
"/dls_sw/apps/zocalo/secrets/credentials-testing.cfg")
StompTransport.add_command_line_options(parser)
parser.add_argument(
"s3_url",
metavar="S3_URL",
help="The access URL for the S3 bucket",
type=ensure_url,
)
# parser.add_argument("images", metavar="IMAGES", help="Image numbers to submit. Index or ranges '1,10' '1-10'. Defauts to all")
parser.add_argument("-v",
"--verbose",
help="increase output verbosity",
action="store_true")
args = parser.parse_args()
help="the DB config file",
default="dbconfig.cfg",
metavar="FILE")
parser.add_option("-s",
"--stomp-config",
dest="stomp_config",
help="the stomp config (i.e. message queue)",
default="stomp.cfg",
metavar="FILE")
(options, args) = parser.parse_args(sys.argv[1:])
config = ConfigParser.RawConfigParser(allow_no_value=True)
config.read(options.config)
set_logging(config)
StompTransport.load_configuration_file(options.stomp_config)
StompTransport.add_command_line_options(parser)
# Get a database connection
mxprocessing = None
with ispyb.open(options.db_config) as conn:
mxprocessing = ispyb.factory.create_data_area(
ispyb.factory.DataAreaType.MXPROCESSING, conn)
def receive_message_but_exit_on_error(*args, **kwargs):
try:
receive_message(*args, **kwargs)
except KeyboardInterrupt:
print("Terminating.")
sys.exit(0)
except Exception as e:
import workflows.recipe
from workflows.transport.stomp_transport import StompTransport
def update_jobid(rw, headers, message):
print('update job id')
print(headers)
if __name__ == '__main__':
default_configuration = '/dls_sw/apps/zocalo/secrets/credentials-live.cfg'
# override default stomp host
try:
StompTransport.load_configuration_file(default_configuration)
except workflows.Error as e:
print("Error: %s\n" % str(e))
# StompTransport.add_command_line_options(parser)
# (options, args) = parser.parse_args(sys.argv[1:])
stomp = StompTransport()
message = {
'recipes': [],
'parameters': {},
}
reply_to = 'transient.scipion.%s' % str(uuid.uuid4())
print('reply to %s' % reply_to)
import os
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.append('/dls_sw/i24/scripts/setup_beamline')
import pv, ca
import Queue
import workflows.recipe
from workflows.transport.stomp_transport import StompTransport
import logging
logging.basicConfig(level=logging.DEBUG)
StompTransport.load_configuration_file(
'/dls_sw/apps/zocalo/secrets/credentials-i24.cfg')
# StompTransport.load_configuration_file('/dls_sw/apps/zocalo/secrets/credentials-live.cfg')
visit_directory = ca.cagetstring(pv.pilat_filepath)
print visit_directory
if visit_directory.startswith('/ramdisk'):
visit_directory = visit_directory.replace('ramdisk', 'dls/i24/data')
print 'visit_director', visit_directory
filefromdet = ca.cagetstring('BL24I-EA-PILAT-01:cam1:FileName_RBV')
pattern = os.path.join(visit_directory, "%s" % (filefromdet) + "%04d.cbf")
print pattern
#pattern = "/dls/i24/data/2018/nt14493-104/had3/agouti/agouti0044_%05d.cbf"
#chip_name = 'chip_name'
#sub_directory='image_analysis'
def test_subscribe_to_broadcast(mockstomp):
"""Test subscribing to a topic (publish-subscribe) and callback functions."""
mock_cb1 = mock.Mock()
mock_cb2 = mock.Mock()
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
def callback_resolver(cbid):
if cbid == 1:
return mock_cb1
if cbid == 2:
return mock_cb2
raise ValueError("Unknown subscription ID %r" % cbid)
stomp.subscription_callback = callback_resolver
mockconn.set_listener.assert_called_once()
listener = mockconn.set_listener.call_args[0][1]
assert listener is not None
stomp._subscribe_broadcast(
1,
str(mock.sentinel.channel1),
mock_cb1,
)
mockconn.subscribe.assert_called_once()
args, kwargs = mockconn.subscribe.call_args
assert args == ("/topic/" + str(mock.sentinel.channel1), 1)
assert kwargs == {"headers": {}}
stomp._subscribe_broadcast(2,
str(mock.sentinel.channel2),
mock_cb2,
retroactive=True)
assert mockconn.subscribe.call_count == 2
args, kwargs = mockconn.subscribe.call_args
assert args == ("/topic/" + str(mock.sentinel.channel2), 2)
assert kwargs == {"headers": {"activemq.retroactive": "true"}}
assert mock_cb1.call_count == 0
listener.on_message(_frame({"subscription": 1}, mock.sentinel.message1))
mock_cb1.assert_called_once_with({"subscription": 1},
mock.sentinel.message1)
assert mock_cb2.call_count == 0
listener.on_message(_frame({"subscription": 2}, mock.sentinel.message2))
mock_cb2.assert_called_once_with({"subscription": 2},
mock.sentinel.message2)
stomp._unsubscribe(1)
mockconn.unsubscribe.assert_called_once_with(id=1)
stomp._unsubscribe(2)
mockconn.unsubscribe.assert_called_with(id=2)
def test_namespace_is_used_correctly(mockstomp):
"""Test that a configured namespace is correctly used when subscribing and sending messages."""
mockconn = mockstomp.Connection.return_value
StompTransport.defaults["--stomp-prfx"] = ""
stomp = StompTransport()
stomp.connect()
assert stomp.get_namespace() == ""
StompTransport.defaults["--stomp-prfx"] = "ns."
stomp = StompTransport()
stomp.connect()
assert stomp.get_namespace() == "ns"
stomp._send("some_queue", mock.sentinel.message1)
mockconn.send.assert_called_once()
assert mockconn.send.call_args[0] == (
"/queue/ns.some_queue",
mock.sentinel.message1,
)
stomp._send("some_queue", mock.sentinel.message2, ignore_namespace=True)
assert mockconn.send.call_args[0] == ("/queue/some_queue",
mock.sentinel.message2)
StompTransport.defaults["--stomp-prfx"] = "ns"
stomp = StompTransport()
stomp.connect()
assert stomp.get_namespace() == "ns"
stomp._send("some_queue", mock.sentinel.message1)
assert mockconn.send.call_args[0] == (
"/queue/ns.some_queue",
mock.sentinel.message1,
)
stomp._broadcast("some_topic", mock.sentinel.message2)
assert mockconn.send.call_args[0] == (
"/topic/ns.some_topic",
mock.sentinel.message2,
)
stomp._broadcast("some_topic",
mock.sentinel.message3,
ignore_namespace=True)
assert mockconn.send.call_args[0] == ("/topic/some_topic",
mock.sentinel.message3)
stomp._subscribe(1, "sub_queue", None)
mockconn.subscribe.assert_called_once()
assert mockconn.subscribe.call_args[0] == ("/queue/ns.sub_queue", 1)
stomp._subscribe(2, "sub_queue", None, ignore_namespace=True)
assert mockconn.subscribe.call_args[0] == ("/queue/sub_queue", 2)
stomp._subscribe_broadcast(3, "sub_topic", None)
assert mockconn.subscribe.call_args[0] == ("/topic/ns.sub_topic", 3)
stomp._subscribe_broadcast(4, "sub_topic", None, ignore_namespace=True)
assert mockconn.subscribe.call_args[0] == ("/topic/sub_topic", 4)
stomp.broadcast_status("some status")
assert mockconn.send.call_args[0] == ("/topic/ns.transient.status",
'"some status"')
def test_transaction_calls(mockstomp):
"""Test that calls to create, commit, abort transactions are passed to stomp properly."""
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
stomp._transaction_begin(mock.sentinel.txid)
mockconn.begin.assert_called_once_with(transaction=mock.sentinel.txid)
stomp._send("destination",
mock.sentinel.message,
transaction=mock.sentinel.txid)
mockconn.send.assert_called_once_with(
"/queue/destination",
mock.sentinel.message,
headers={"persistent": "true"},
transaction=mock.sentinel.txid,
)
stomp._transaction_abort(mock.sentinel.txid)
mockconn.abort.assert_called_once_with(mock.sentinel.txid)
stomp._transaction_commit(mock.sentinel.txid)
mockconn.commit.assert_called_once_with(mock.sentinel.txid)