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],