def test_execute_different_project_ids(self, mock_hook):
another_project = 'another-project'
operator = PubSubCreateSubscriptionOperator(
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 = PubSubCreateSubscriptionOperator(task_id=TASK_ID,
project_id=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(
project_id=TEST_PROJECT,
topic=TEST_TOPIC,
subscription=None,
subscription_project_id=None,
ack_deadline_secs=10,
fail_if_exists=False,
push_config=None,
retain_acked_messages=None,
message_retention_duration=None,
labels=None,
enable_message_ordering=False,
expiration_policy=None,
filter_=None,
dead_letter_policy=None,
retry_policy=None,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(response, TEST_SUBSCRIPTION)
with models.DAG(
"example_gcp_pubsub_sensor",
schedule_interval='@once', # Override to match your needs
start_date=START_DATE,
catchup=False,
) as example_sensor_dag:
# [START howto_operator_gcp_pubsub_create_topic]
create_topic = PubSubCreateTopicOperator(task_id="create_topic",
topic=TOPIC_FOR_SENSOR_DAG,
project_id=GCP_PROJECT_ID,
fail_if_exists=False)
# [END howto_operator_gcp_pubsub_create_topic]
# [START howto_operator_gcp_pubsub_create_subscription]
subscribe_task = PubSubCreateSubscriptionOperator(
task_id="subscribe_task",
project_id=GCP_PROJECT_ID,
topic=TOPIC_FOR_SENSOR_DAG)
# [END howto_operator_gcp_pubsub_create_subscription]
# [START howto_operator_gcp_pubsub_pull_message_with_sensor]
subscription = subscribe_task.output
pull_messages = PubSubPullSensor(
task_id="pull_messages",
ack_messages=True,
project_id=GCP_PROJECT_ID,
subscription=subscription,
)
# [END howto_operator_gcp_pubsub_pull_message_with_sensor]
# [START howto_operator_gcp_pubsub_pull_messages_result]
# [END howto_operator_gcp_pubsub_pull_messages_result_cmd]
with models.DAG(
"example_gcp_pubsub",
default_args=default_args,
schedule_interval=None, # Override to match your needs
) as example_dag:
# [START howto_operator_gcp_pubsub_create_topic]
create_topic = PubSubCreateTopicOperator(
task_id="create_topic", topic=TOPIC, project_id=GCP_PROJECT_ID
)
# [END howto_operator_gcp_pubsub_create_topic]
# [START howto_operator_gcp_pubsub_create_subscription]
subscribe_task = PubSubCreateSubscriptionOperator(
task_id="subscribe_task", project_id=GCP_PROJECT_ID, topic=TOPIC
)
# [END howto_operator_gcp_pubsub_create_subscription]
# [START howto_operator_gcp_pubsub_pull_message]
subscription = "{{ task_instance.xcom_pull('subscribe_task') }}"
pull_messages = PubSubPullSensor(
task_id="pull_messages",
ack_messages=True,
project_id=GCP_PROJECT_ID,
subscription=subscription,
)
# [END howto_operator_gcp_pubsub_pull_message]
# [START howto_operator_gcp_pubsub_pull_messages_result]