def test_execute_timeout(self, mock_hook): operator = PubSubPullSensor(task_id=TASK_ID, project=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_with( TEST_PROJECT, TEST_SUBSCRIPTION, 5, False)
def test_execute(self, mock_hook): operator = PubSubPullSensor(task_id=TASK_ID, project=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, poke_interval=0) generated_messages = self._generate_messages(5) mock_hook.return_value.pull.return_value = generated_messages response = operator.execute(None) mock_hook.return_value.pull.assert_called_with( TEST_PROJECT, TEST_SUBSCRIPTION, 5, False) self.assertEquals(response, generated_messages)
def test_poke_with_ack_messages(self, mock_hook): operator = PubSubPullSensor(task_id=TASK_ID, project=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, ack_messages=True) generated_messages = self._generate_messages(5) mock_hook.return_value.pull.return_value = generated_messages self.assertEquals(generated_messages, operator.poke(None)) mock_hook.return_value.acknowledge.assert_called_with( TEST_PROJECT, TEST_SUBSCRIPTION, ['1', '2', '3', '4', '5'] )
'email_on_failure': False, 'email_on_retry': False, 'project': project, 'topic': topic, 'subscription': subscription, } echo_template = ''' {% for m in task_instance.xcom_pull(task_ids='pull-messages') %} echo "AckID: {{ m.get('ackId') }}, Base64-Encoded: {{ m.get('message') }}" {% endfor %} ''' with DAG('pubsub-end-to-end', default_args=default_args, schedule_interval=datetime.timedelta(days=1)) as dag: t1 = PubSubTopicCreateOperator(task_id='create-topic') t2 = PubSubSubscriptionCreateOperator( task_id='create-subscription', topic_project=project, subscription=subscription) t3 = PubSubPublishOperator( task_id='publish-messages', messages=messages) t4 = PubSubPullSensor(task_id='pull-messages', ack_messages=True) t5 = BashOperator(task_id='echo-pulled-messages', bash_command=echo_template) t6 = PubSubSubscriptionDeleteOperator(task_id='delete-subscription') t7 = PubSubTopicDeleteOperator(task_id='delete-topic') t1 >> t2 >> t3 t2 >> t4 >> t5 >> t6 >> t7
def test_poke_no_messages(self, mock_hook): operator = PubSubPullSensor(task_id=TASK_ID, project=TEST_PROJECT, subscription=TEST_SUBSCRIPTION) mock_hook.return_value.pull.return_value = [] self.assertEquals([], operator.poke(None))