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 = mock.MagicMock()
self.sqs_hook = SQSHook()
# Using a task-decorated function to create an SQS queue
create_queue = create_queue_fn()
publish_to_queue = SQSPublishOperator(
task_id='publish_to_queue',
sqs_queue=create_queue,
message_content="{{ task_instance }}-{{ execution_date }}",
message_attributes=None,
delay_seconds=0,
aws_conn_id=AWS_CONN_ID,
)
read_from_queue = SQSSensor(
task_id='read_from_queue',
sqs_queue=create_queue,
max_messages=5,
wait_time_seconds=1,
visibility_timeout=None,
message_filtering=None,
message_filtering_match_values=None,
message_filtering_config=None,
aws_conn_id=AWS_CONN_ID,
)
# Using a task-decorated function to delete the SQS queue we created earlier
delete_queue = delete_queue_fn(create_queue)
create_queue >> publish_to_queue >> read_from_queue >> delete_queue
# [END howto_sqs_operator_and_sensor]
class TestSQSSensor(unittest.TestCase):
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 = mock.MagicMock()
self.sqs_hook = SQSHook()
@mock_sqs
def test_poke_success(self):
self.sqs_hook.create_queue('test')
self.sqs_hook.send_message(queue_url='test', message_body='hello')
result = self.sensor.poke(self.mock_context)
self.assertTrue(result)
self.assertTrue(
"'Body': 'hello'" in str(self.mock_context['ti'].method_calls),
"context call should contain message hello",
)
@mock_sqs
def test_poke_no_messsage_failed(self):
self.sqs_hook.create_queue('test')
result = self.sensor.poke(self.mock_context)
self.assertFalse(result)
context_calls = []
self.assertTrue(self.mock_context['ti'].method_calls == context_calls,
"context call should be same")
@mock.patch.object(SQSHook, 'get_conn')
def test_poke_delete_raise_airflow_exception(self, mock_conn):
message = {
'Messages': [{
'MessageId': 'c585e508-2ea0-44c7-bf3e-d1ba0cb87834',
'ReceiptHandle': 'mockHandle',
'MD5OfBody': 'e5a9d8684a8edfed460b8d42fd28842f',
'Body': 'h21',
}],
'ResponseMetadata': {
'RequestId': '56cbf4aa-f4ef-5518-9574-a04e0a5f1411',
'HTTPStatusCode': 200,
'HTTPHeaders': {
'x-amzn-requestid': '56cbf4aa-f4ef-5518-9574-a04e0a5f1411',
'date': 'Mon, 18 Feb 2019 18:41:52 GMT',
'content-type': 'text/xml',
'mock_sqs_hook-length': '830',
},
'RetryAttempts': 0,
},
}
mock_conn.return_value.receive_message.return_value = message
mock_conn.return_value.delete_message_batch.return_value = {
'Failed': [{
'Id': '22f67273-4dbc-4c19-83b5-aee71bfeb832'
}]
}
with self.assertRaises(AirflowException) as context:
self.sensor.poke(self.mock_context)
self.assertTrue(
'Delete SQS Messages failed' in context.exception.args[0])
@mock.patch.object(SQSHook, 'get_conn')
def test_poke_receive_raise_exception(self, mock_conn):
mock_conn.return_value.receive_message.side_effect = Exception(
'test exception')
with self.assertRaises(Exception) as context:
self.sensor.poke(self.mock_context)
self.assertTrue('test exception' in context.exception.args[0])