def test_listen_for_messages(self, mock_fetch_and_process, mock_get_queue):
num_messages = 3
visibility_timeout_s = 4
loop_count = 1
listen_for_messages(Priority.high, num_messages, visibility_timeout_s, loop_count)
queue_name = get_queue_name(Priority.high)
mock_get_queue.assert_called_once_with(queue_name)
mock_fetch_and_process.assert_called_once_with(
queue_name, mock_get_queue.return_value, num_messages=num_messages, visibility_timeout=visibility_timeout_s
)
def extend_visibility_timeout(priority: Priority, receipt: str,
visibility_timeout_s: int) -> None:
"""
Extends visibility timeout of a message on a given priority queue for long running tasks.
"""
queue_name = get_queue_name(priority)
client = _get_sqs_client()
queue_url = _get_queue_url(client, queue_name)
client.change_message_visibility(QueueUrl=queue_url,
ReceiptHandle=receipt,
VisibilityTimeout=visibility_timeout_s)
def publish(message: Message) -> None:
"""
Publishes a message on Taskhawk queue
"""
message_body = message.as_dict()
payload = _convert_to_json(message_body)
if settings.IS_LAMBDA_APP:
topic = _get_sns_topic(message.priority)
_publish_over_sns(topic, payload, message.headers)
else:
queue_name = get_queue_name(message.priority)
queue = get_queue(queue_name)
_publish_over_sqs(queue, payload, message.headers)
_log_published_message(message_body)
def test_requeue_dead_letter(mock_get_queue, mock_get_queue_messages):
priority = Priority.high
num_messages = 3
visibility_timeout = 4
messages = [mock.MagicMock() for _ in range(num_messages)]
mock_get_queue_messages.side_effect = iter([messages, None])
dlq_name = f'{get_queue_name(priority)}-DLQ'
mock_queue, mock_dlq = mock.MagicMock(), mock.MagicMock()
mock_queue.attributes = {
'RedrivePolicy': json.dumps({'deadLetterTargetArn': dlq_name})
}
mock_queue.send_messages.return_value = {'Failed': []}
mock_get_queue.side_effect = iter([mock_queue, mock_dlq])
mock_dlq.delete_messages.return_value = {'Failed': []}
requeue_dead_letter(priority,
num_messages=num_messages,
visibility_timeout=visibility_timeout)
mock_get_queue.assert_has_calls(
[mock.call(get_queue_name(priority)),
mock.call(dlq_name)])
mock_get_queue_messages.assert_has_calls([
mock.call(mock_dlq,
num_messages=num_messages,
visibility_timeout=visibility_timeout),
mock.call(mock_dlq,
num_messages=num_messages,
visibility_timeout=visibility_timeout),
])
mock_queue.send_messages.assert_called_once_with(
Entries=[{
'Id': queue_message.message_id,
'MessageBody': queue_message.body,
'MessageAttributes': queue_message.message_attributes,
} for queue_message in messages])
mock_dlq.delete_messages.assert_called_once_with(
Entries=[{
'Id': queue_message.message_id,
'ReceiptHandle': queue_message.receipt_handle
} for queue_message in messages])
def test_requeue_dead_letter_failure(mock_get_queue, mock_get_queue_messages):
priority = Priority.high
num_messages = 3
visibility_timeout = 4
messages = [mock.MagicMock() for _ in range(num_messages)]
mock_get_queue_messages.side_effect = iter([messages, None])
dlq_name = f'{get_queue_name(priority)}-DLQ'
mock_queue, mock_dlq = mock.MagicMock(), mock.MagicMock()
mock_queue.attributes = {
'RedrivePolicy': json.dumps({'deadLetterTargetArn': dlq_name})
}
mock_queue.send_messages.return_value = {
'Failed': [{
'Id': 'string'
}],
'Successful': []
}
mock_get_queue.side_effect = iter([mock_queue, mock_dlq])
with pytest.raises(PartialFailure) as exc_info:
requeue_dead_letter(priority, num_messages, visibility_timeout)
assert exc_info.value.success_count == 0
assert exc_info.value.failure_count == 1
mock_get_queue.assert_has_calls(
[mock.call(get_queue_name(priority)),
mock.call(dlq_name)])
# not called a 2nd time after failure
mock_get_queue_messages.assert_called_once_with(
mock_dlq,
num_messages=num_messages,
visibility_timeout=visibility_timeout)
mock_queue.send_messages.assert_called_once_with(
Entries=[{
'Id': queue_message.message_id,
'MessageBody': queue_message.body,
'MessageAttributes': queue_message.message_attributes,
} for queue_message in messages])
mock_dlq.delete_messages.assert_not_called()
def requeue_dead_letter(priority: Priority,
num_messages: int = 10,
visibility_timeout: int = None) -> None:
"""
Re-queues everything in the Taskhawk DLQ back into the Taskhawk queue.
:param priority: The priority queue to listen to
:param num_messages: Maximum number of messages to fetch in one SQS call. Defaults to 10.
:param visibility_timeout: The number of seconds the message should remain invisible to other queue readers.
Defaults to None, which is queue default
"""
if settings.IS_LAMBDA_APP:
logging.warning("Not supported for Lambda apps")
return
queue = get_queue(get_queue_name(priority))
dead_letter_queue = get_dead_letter_queue(queue)
logging.info("Re-queueing messages from {} to {}".format(
dead_letter_queue.url, queue.url))
while True:
queue_messages = get_queue_messages(
dead_letter_queue,
num_messages=num_messages,
visibility_timeout=visibility_timeout)
if not queue_messages:
break
logging.info("got {} messages from dlq".format(len(queue_messages)))
_enqueue_messages(queue, queue_messages)
response = dead_letter_queue.delete_messages(
Entries=[{
'Id': message.message_id,
'Receipt': message.receipt_handle
} for message in queue_messages])
if response['Failed']:
raise PartialFailure(response)
logging.info("Re-queued {} messages".format(len(queue_messages)))
def test_get_queue_name(priority, suffix):
assert get_queue_name(
priority) == f'TASKHAWK-{settings.TASKHAWK_QUEUE.upper()}{suffix}'
