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 = PubSubPullSensor( task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, poke_interval=0, 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_poke_no_messages(self, mock_hook): operator = PubSubPullSensor( task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, ) mock_hook.return_value.pull.return_value = [] self.assertEqual(False, operator.poke({}))
def test_poke_with_ack_messages(self, mock_hook): operator = PubSubPullSensor(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.poke(None)) mock_hook.return_value.acknowledge.assert_called_once_with( project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, ack_ids=['1', '2', '3', '4', '5'])
def test_execute_timeout(self, mock_hook): operator = PubSubPullSensor(task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, poke_interval=0, timeout=1) mock_hook.return_value.pull.return_value = [] with self.assertRaises(AirflowSensorTimeout): operator.execute(None) mock_hook.return_value.pull.assert_called_once_with( project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=5, return_immediately=False )
def test_poke_with_ack_messages(self, mock_hook): operator = PubSubPullSensor( task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, ack_messages=True, ) generated_messages = self._generate_messages(5) mock_hook.return_value.pull.return_value = generated_messages self.assertEqual(True, operator.poke({})) mock_hook.return_value.acknowledge.assert_called_once_with( project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, messages=generated_messages, )
def test_execute(self, mock_hook): operator = PubSubPullSensor( task_id=TASK_ID, project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, poke_interval=0, ) generated_messages = self._generate_messages(5) generated_dicts = self._generate_dicts(5) 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 ) self.assertEqual(generated_dicts, response)
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] 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_SENSOR_DAG, messages=[MESSAGE] * 10,
default_args=DEFAULT_DAG_ARGS) subscription = 'latihan-subs' gcp_conn_id = 'google_cloud_default' project_id = 'sekolah-bigdata-3' #define task task_start = BashOperator( task_id="start_task", bash_command='echo start', dag=dag ) pull_message = PubSubPullSensor( task_id="pull_message_pubsub", ack_messages=True, project_id=project_id, gcp_conn_id=gcp_conn_id, subscription=subscription, dag=dag ) task_end = BashOperator( task_id="end_task", bash_command='echo end', dag=dag ) task_start >> pull_message >> task_end