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