def setUp(self):
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.sensor = SQSSensor(task_id='test_task',
dag=self.dag,
sqs_queue='test',
aws_conn_id='aws_default')
self.mock_context = MagicMock()
self.sqs_hook = SQSHook()
def execute(self, context):
"""
Publish the message to SQS queue
:param context: the context object
:type context: dict
:return: dict with information about the message sent
For details of the returned dict see :py:meth:`botocore.client.SQS.send_message`
:rtype: dict
"""
hook = SQSHook(aws_conn_id=self.aws_conn_id)
result = hook.send_message(queue_url=self.sqs_queue,
message_body=self.message_content,
delay_seconds=self.delay_seconds,
message_attributes=self.message_attributes)
self.log.info('result is send_message is %s', result)
return result
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("received message %s", str(messages))
if 'Messages' in messages and messages['Messages']:
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())