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 create_queue_fn(): """This is a python callback that creates an SQS queue""" hook = SQSHook(aws_conn_id=AWS_CONN_ID) result = hook.create_queue(queue_name=QUEUE_NAME, ) return result['QueueUrl']
def create_queue_fn(): """This is a Python function that creates an SQS queue""" hook = SQSHook() result = hook.create_queue(queue_name=QUEUE_NAME) return result['QueueUrl']
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])