class TestSQSPublishOperator(unittest.TestCase):
def setUp(self):
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.operator = SQSPublishOperator(task_id='test_task',
dag=self.dag,
sqs_queue='test',
message_content='hello',
aws_conn_id='aws_default')
self.mock_context = MagicMock()
self.sqs_hook = SQSHook()
@mock_sqs
def test_execute_success(self):
self.sqs_hook.create_queue('test')
result = self.operator.execute(self.mock_context)
self.assertTrue('MD5OfMessageBody' in result)
self.assertTrue('MessageId' in result)
message = self.sqs_hook.get_conn().receive_message(QueueUrl='test')
self.assertEqual(len(message['Messages']), 1)
self.assertEqual(message['Messages'][0]['MessageId'],
result['MessageId'])
self.assertEqual(message['Messages'][0]['Body'], 'hello')
context_calls = []
self.assertTrue(self.mock_context['ti'].method_calls == context_calls,
"context call should be same")
def poke(self, context):
"""
Check for message on subscribed queue and write to xcom the message with key ``messages``
:param context: the context object
:type context: dict
:return: ``True`` if message is available or ``False``
"""
sqs_hook = SQSHook(aws_conn_id=self.aws_conn_id)
sqs_conn = sqs_hook.get_conn()
self.log.info('SQSSensor checking for message on queue: %s',
self.sqs_queue)
messages = sqs_conn.receive_message(
QueueUrl=self.sqs_queue,
MaxNumberOfMessages=self.max_messages,
WaitTimeSeconds=self.wait_time_seconds)
self.log.info("reveived message %s", str(messages))
if 'Messages' in messages and len(messages['Messages']) > 0:
entries = [{
'Id': message['MessageId'],
'ReceiptHandle': message['ReceiptHandle']
} for message in messages['Messages']]
result = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue,
Entries=entries)
if 'Successful' in result:
context['ti'].xcom_push(key='messages', value=messages)
return True
else:
raise AirflowException('Delete SQS Messages failed ' +
str(result) + ' for messages ' +
str(messages))
return False
def poke(self, context):
"""
Check for message on subscribed queue and write to xcom the message with key ``messages``
:param context: the context object
:type context: dict
:return: ``True`` if message is available or ``False``
"""
sqs_hook = SQSHook(aws_conn_id=self.aws_conn_id)
sqs_conn = sqs_hook.get_conn()
self.log.info('SQSSensor checking for message on queue: %s', self.sqs_queue)
messages = sqs_conn.receive_message(QueueUrl=self.sqs_queue,
MaxNumberOfMessages=self.max_messages,
WaitTimeSeconds=self.wait_time_seconds)
self.log.info("reveived message %s", str(messages))
if 'Messages' in messages and len(messages['Messages']) > 0:
entries = [{'Id': message['MessageId'], 'ReceiptHandle': message['ReceiptHandle']}
for message in messages['Messages']]
result = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue,
Entries=entries)
if 'Successful' in result:
context['ti'].xcom_push(key='messages', value=messages)
return True
else:
raise AirflowException(
'Delete SQS Messages failed ' + str(result) + ' for messages ' + str(messages))
return False
def test_get_conn(self):
hook = SQSHook(aws_conn_id='aws_default')
self.assertIsNotNone(hook.get_conn())
def test_get_conn(self):
hook = SQSHook(aws_conn_id='aws_default')
self.assertIsNotNone(hook.get_conn())