def test_execute_different_project_ids(self, mock_hook):
another_project = 'another-project'
operator = PubSubSubscriptionCreateOperator(
project_id=TEST_PROJECT,
topic=TEST_TOPIC,
subscription=TEST_SUBSCRIPTION,
subscription_project_id=another_project,
task_id=TASK_ID
)
mock_hook.return_value.create_subscription.return_value = TEST_SUBSCRIPTION
response = operator.execute(None)
mock_hook.return_value.create_subscription.assert_called_once_with(
project_id=TEST_PROJECT,
topic=TEST_TOPIC,
subscription=TEST_SUBSCRIPTION,
subscription_project_id=another_project,
ack_deadline_secs=10,
fail_if_exists=False,
push_config=None,
retain_acked_messages=None,
message_retention_duration=None,
labels=None,
retry=None,
timeout=None,
metadata=None
)
self.assertEqual(response, TEST_SUBSCRIPTION)
def test_execute_no_subscription(self, mock_hook):
operator = PubSubSubscriptionCreateOperator(task_id=TASK_ID,
topic_project=TEST_PROJECT,
topic=TEST_TOPIC)
mock_hook.return_value.create_subscription.return_value = (
TEST_SUBSCRIPTION)
response = operator.execute(None)
mock_hook.return_value.create_subscription.assert_called_once_with(
TEST_PROJECT, TEST_TOPIC, None, None, 10, False)
self.assertEqual(response, TEST_SUBSCRIPTION)
def test_execute_different_project_ids(self, mock_hook):
another_project = 'another-project'
operator = PubSubSubscriptionCreateOperator(
task_id=TASK_ID,
topic_project=TEST_PROJECT,
topic=TEST_TOPIC,
subscription=TEST_SUBSCRIPTION,
subscription_project=another_project)
mock_hook.return_value.create_subscription.return_value = (
TEST_SUBSCRIPTION)
response = operator.execute(None)
mock_hook.return_value.create_subscription.assert_called_once_with(
TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION, another_project, 10,
False)
self.assertEqual(response, TEST_SUBSCRIPTION)
TOPIC = "PubSubTestTopic"
MESSAGE = {"attributes": {"name": "wrench", "mass": "1.3kg", "count": "3"}}
default_args = {"start_date": airflow.utils.dates.days_ago(1)}
with models.DAG(
"example_gcp_pubsub",
default_args=default_args,
schedule_interval=None, # Override to match your needs
) as example_dag:
create_topic = PubSubTopicCreateOperator(
task_id="create_topic", topic=TOPIC, project=GCP_PROJECT_ID
)
subscribe_task = PubSubSubscriptionCreateOperator(
task_id="subscribe_task", topic_project=GCP_PROJECT_ID, topic=TOPIC
)
publish_task = PubSubPublishOperator(
task_id="publish_task", project=GCP_PROJECT_ID, topic=TOPIC, messages=[MESSAGE]
)
unsubscribe_task = PubSubSubscriptionDeleteOperator(
task_id="unsubscribe_task",
project=GCP_PROJECT_ID,
subscription="{{ task_instance.xcom_pull('subscribe_task') }}",
)
delete_topic = PubSubTopicDeleteOperator(task_id="delete_topic", topic=TOPIC, project=GCP_PROJECT_ID)
create_topic >> [subscribe_task, publish_task] >> unsubscribe_task >> delete_topic
'email_on_retry': False,
'project': project,
'topic': topic,
'subscription': subscription,
}
echo_template = '''
{% for m in task_instance.xcom_pull(task_ids='pull-messages') %}
echo "AckID: {{ m.get('ackId') }}, Base64-Encoded: {{ m.get('message') }}"
{% endfor %}
'''
with DAG('pubsub-end-to-end',
default_args=default_args,
schedule_interval=datetime.timedelta(days=1)) as dag:
# pylint: disable=no-value-for-parameter
# Note that parameters gets passed via dags default_args above.
t1 = PubSubTopicCreateOperator(task_id='create-topic')
t2 = PubSubSubscriptionCreateOperator(task_id='create-subscription',
topic_project=project,
subscription=subscription)
t3 = PubSubPublishOperator(task_id='publish-messages', messages=messages)
t4 = PubSubPullSensor(task_id='pull-messages', ack_messages=True)
t5 = BashOperator(task_id='echo-pulled-messages',
bash_command=echo_template)
t6 = PubSubSubscriptionDeleteOperator(task_id='delete-subscription')
t7 = PubSubTopicDeleteOperator(task_id='delete-topic')
t1 >> t2 >> t3
t2 >> t4 >> t5 >> t6 >> t7
