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
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_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
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_instantiate_link_and_connect_to_broker(mockstomp):
"""Test the Stomp connection routine."""
stomp = StompTransport()
mockconn = mockstomp.Connection.return_value
assert not stomp.is_connected()
stomp.connect()
mockstomp.Connection.assert_called_once()
mockconn.connect.assert_called_once()
assert stomp.is_connected()
stomp.connect()
mockstomp.Connection.assert_called_once()
mockconn.connect.assert_called_once()
assert stomp.is_connected()
stomp.disconnect()
mockstomp.Connection.assert_called_once()
mockconn.connect.assert_called_once()
mockconn.disconnect.assert_called_once()
assert not stomp.is_connected()
stomp.disconnect()
mockstomp.Connection.assert_called_once()
mockconn.connect.assert_called_once()
mockconn.disconnect.assert_called_once()
assert not stomp.is_connected()
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 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_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_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 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 test_messages_are_deserialized_after_transport(mockstomp):
"""Test the message serialization."""
banana = {"entry": [0, "banana"]}
banana_str = '{"entry": [0, "banana"]}'
stomp = StompTransport()
stomp.connect()
mockconn = mockstomp.Connection.return_value
message_handler = mockconn.set_listener.call_args[0][1].on_message
# Test subscriptions
callback = mock.Mock()
stomp.subscribe("channel", callback)
subscription_id = mockconn.subscribe.call_args[0][1]
message_handler(_frame({"subscription": subscription_id}, banana_str))
callback.assert_called_once_with({"subscription": subscription_id}, banana)
message_handler(
_frame({"subscription": subscription_id},
mock.sentinel.undeserializable))
callback.assert_called_with({"subscription": subscription_id},
mock.sentinel.undeserializable)
# Test broadcast subscriptions
callback = mock.Mock()
stomp.subscribe_broadcast("channel", callback)
subscription_id = mockconn.subscribe.call_args[0][1]
message_handler(_frame({"subscription": subscription_id}, banana_str))
callback.assert_called_once_with({"subscription": subscription_id}, banana)
message_handler(
_frame({"subscription": subscription_id},
mock.sentinel.undeserializable))
callback.assert_called_with({"subscription": subscription_id},
mock.sentinel.undeserializable)
# Test subscriptions with mangling disabled
callback = mock.Mock()
stomp.subscribe("channel", callback, disable_mangling=True)
subscription_id = mockconn.subscribe.call_args[0][1]
message_handler(_frame({"subscription": subscription_id}, banana_str))
callback.assert_called_once_with({"subscription": subscription_id},
banana_str)
# Test broadcast subscriptions with mangling disabled
callback = mock.Mock()
stomp.subscribe_broadcast("channel", callback, disable_mangling=True)
subscription_id = mockconn.subscribe.call_args[0][1]
message_handler(_frame({"subscription": subscription_id}, banana_str))
callback.assert_called_once_with({"subscription": subscription_id},
banana_str)
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_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_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)
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_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_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 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_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)
recipe['1']['parameters'] = {}
recipe['1']['parameters']['arguments'] = sys.argv[1:] + ['>', output_file]
recipe['1']['parameters']['cwd'] = os.getcwd()
message['custom_recipe'] = recipe
# reply_to = 'transient.gctf.%s' % str(uuid.uuid4())
# recipe['1']['output'] = 2
# recipe['2'] = {}
# recipe['2']['service'] = "relion_refine_call_back"
# recipe['2']['queue'] = reply_to
# recipe['2']['parameters'] = {}
# recipe['2']['output'] = 3
# recipe['3'] = {}
recipe['start'] = [[1, []]]
stomp.connect()
test_valid_recipe = workflows.recipe.Recipe(recipe)
test_valid_recipe.validate()
stomp.send('processing_recipe', message)
# So that BASH can pick this up
# get_output_file()
print(output_file)
# print("\nMotioncor2 job submitted")
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_subscribe_to_queue(mockstomp):
"""Test subscribing to a queue (producer-consumer), callback functions and unsubscribe."""
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(
1,
str(mock.sentinel.channel1),
mock_cb1,
)
mockconn.subscribe.assert_called_once()
args, kwargs = mockconn.subscribe.call_args
assert args == ("/queue/" + str(mock.sentinel.channel1), 1)
assert kwargs == {
"headers": {},
"ack": "auto",
}
stomp._subscribe(
2,
str(mock.sentinel.channel2),
mock_cb2,
retroactive=True,
selector=mock.sentinel.selector,
exclusive=True,
priority=42,
)
assert mockconn.subscribe.call_count == 2
args, kwargs = mockconn.subscribe.call_args
assert args == ("/queue/" + str(mock.sentinel.channel2), 2)
assert kwargs == {
"headers": {
"activemq.retroactive": "true",
"selector": mock.sentinel.selector,
"activemq.exclusive": "true",
"activemq.priority": 42,
},
"ack": "auto",
}
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._subscribe(3,
str(mock.sentinel.channel3),
mock_cb2,
acknowledgement=True)
assert mockconn.subscribe.call_count == 3
args, kwargs = mockconn.subscribe.call_args
assert args == ("/queue/" + str(mock.sentinel.channel3), 3)
assert kwargs == {"headers": {}, "ack": "client-individual"}
stomp._unsubscribe(1)
mockconn.unsubscribe.assert_called_once_with(id=1)
stomp._unsubscribe(2)
mockconn.unsubscribe.assert_called_with(id=2)
