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_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']
     )
 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)
    '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))