def test_execute(self, mock_hook):
# Given
hook_response = {'MessageId': 'foobar'}
hook_instance = mock_hook.return_value
hook_instance.publish_to_target.return_value = hook_response
operator = SnsPublishOperator(task_id=TASK_ID,
aws_conn_id=AWS_CONN_ID,
target_arn=TARGET_ARN,
message=MESSAGE)
# When
result = operator.execute(None)
# Then
self.assertEqual(hook_response, result)
def test_execute(self, mock_hook):
# Given
hook_response = {'MessageId': 'foobar'}
hook_instance = mock_hook.return_value
hook_instance.publish_to_target.return_value = hook_response
operator = SnsPublishOperator(
task_id=TASK_ID,
aws_conn_id=AWS_CONN_ID,
target_arn=TARGET_ARN,
message=MESSAGE
)
# When
result = operator.execute(None)
# Then
self.assertEqual(hook_response, result)
def test_init(self):
# Given / When
operator = SnsPublishOperator(task_id=TASK_ID,
aws_conn_id=AWS_CONN_ID,
target_arn=TARGET_ARN,
message=MESSAGE)
# Then
self.assertEqual(TASK_ID, operator.task_id)
self.assertEqual(AWS_CONN_ID, operator.aws_conn_id)
self.assertEqual(TARGET_ARN, operator.target_arn)
self.assertEqual(MESSAGE, operator.message)
def test_init(self):
# Given / When
operator = SnsPublishOperator(
task_id=TASK_ID,
aws_conn_id=AWS_CONN_ID,
target_arn=TARGET_ARN,
message=MESSAGE,
subject=SUBJECT,
message_attributes=MESSAGE_ATTRIBUTES,
)
# Then
self.assertEqual(TASK_ID, operator.task_id)
self.assertEqual(AWS_CONN_ID, operator.aws_conn_id)
self.assertEqual(TARGET_ARN, operator.target_arn)
self.assertEqual(MESSAGE, operator.message)
self.assertEqual(SUBJECT, operator.subject)
self.assertEqual(MESSAGE_ATTRIBUTES, operator.message_attributes)
def _failure_callback(context: Mapping[str, Any]):
"""Callback for when a task fails."""
dag: PMIDAG = context['dag']
task_instance: TaskInstance = context['ti']
return SnsPublishOperator(
task_id='task_failure_callback',
aws_conn_id=dag.aws_conn_id,
target_arn=':'.join([
"arn:aws:sns", dag.aws_env['region'], dag.aws_env['accountId'],
dag.env_config['alarms_topic_name']
]),
subject='PMI DAG run failure',
message=(
f"Airflow task failure:\n"
f"\n"
f"\nDAG: {task_instance.dag_id}"
f"\nTask: {task_instance.task_id}"
f"\nExecution Date: {task_instance.execution_date}"
f"\nHost: {task_instance.hostname}"
f"\nError: {context['exception']}")).execute(context)
python_callable=get_stale_partition,
provide_context=True,
dag=dag
)
demo_check_operator = BranchPythonOperator(
task_id='demo_check',
python_callable=demo_check,
trigger_rule='one_success',
dag=dag
)
sns_operator = SnsPublishOperator(
task_id='send_sns',
target_arn='arn:aws:sns:us-west-2:356032320829:airflow',
message='Airflow fingerprinting done',
aws_conn_id='S3_conn',
trigger_rule='one_success',
dag=dag
)
#################### Flow ###############
dummy_operator >> get_stale_partition_operator >> demo_check_operator
demo_check_operator >> sns_operator
demo_check_operator >> skip_operator
sns_operator >> end_operator
skip_operator >> end_operator
hook = AwsSnsHook()
content_object = hook.get_key(key=key, bucket_name=bucket_name)
file_content = content_object.get()['Body'].read().decode('utf-8')
return json.loads(file_content)
"""
# Sends a test message from Amazon SNS from Amazon MWAA.
"""
with DAG(
dag_id=DAG_ID,
description='Send a test message to an SNS Topic',
default_args=DEFAULT_ARGS,
dagrun_timeout=timedelta(hours=2),
start_date=days_ago(1),
schedule_interval=None,
tags=['sns']
) as dag:
dag.doc_md = __doc__
sns_publish = SnsPublishOperator(
task_id='publish_sns_message',
target_arn='{{ var.value.sns_topic }}',
message='Test message sent to Amazon SNS from Amazon MWAA.',
aws_conn_id='aws_default',
subject='Test Message from Amazon MWAA'
)
sns_publish