def test_execute_with_messages_callback(self, mock_hook):
generated_messages = self._generate_messages(5)
messages_callback_return_value = 'asdfg'
def messages_callback(
pulled_messages: List[ReceivedMessage],
context: Dict[str, Any],
):
assert pulled_messages == generated_messages
assert isinstance(context, dict)
for key in context.keys():
assert isinstance(key, str)
return messages_callback_return_value
messages_callback = mock.Mock(side_effect=messages_callback)
operator = PubSubPullOperator(
task_id=TASK_ID,
project_id=TEST_PROJECT,
subscription=TEST_SUBSCRIPTION,
messages_callback=messages_callback,
)
mock_hook.return_value.pull.return_value = generated_messages
response = operator.execute({})
mock_hook.return_value.pull.assert_called_once_with(
project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=5, return_immediately=True
)
messages_callback.assert_called_once()
assert response == messages_callback_return_value
def test_execute_no_messages(self, mock_hook):
operator = PubSubPullOperator(
task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION,
)
mock_hook.return_value.pull.return_value = []
self.assertEqual([], operator.execute({}))
def test_execute_with_ack_messages(self, mock_hook):
operator = PubSubPullOperator(
task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, ack_messages=True,
)
generated_messages = self._generate_messages(5)
generated_dicts = self._generate_dicts(5)
mock_hook.return_value.pull.return_value = generated_messages
self.assertEqual(generated_dicts, operator.execute({}))
mock_hook.return_value.acknowledge.assert_called_once_with(
project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, messages=generated_messages,
)
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_FOR_OPERATOR_DAG)
# [END howto_operator_gcp_pubsub_create_subscription]
# [START howto_operator_gcp_pubsub_pull_message_with_operator]
subscription = subscribe_task.output
pull_messages_operator = PubSubPullOperator(
task_id="pull_messages",
ack_messages=True,
project_id=GCP_PROJECT_ID,
subscription=subscription,
)
# [END howto_operator_gcp_pubsub_pull_message_with_operator]
# [START howto_operator_gcp_pubsub_pull_messages_result]
pull_messages_result = BashOperator(task_id="pull_messages_result",
bash_command=echo_cmd)
# [END howto_operator_gcp_pubsub_pull_messages_result]
# [START howto_operator_gcp_pubsub_publish]
publish_task = PubSubPublishMessageOperator(
task_id="publish_task",
project_id=GCP_PROJECT_ID,
topic=TOPIC_FOR_OPERATOR_DAG,
messages=[MESSAGE, MESSAGE, MESSAGE],
