class PubSubHookTest(unittest.TestCase):
    def setUp(self):
        with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
                        new=mock_init):
            self.pubsub_hook = PubSubHook(gcp_conn_id='test')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_topic(self, mock_service):
        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

        create_method = (mock_service.return_value.projects.return_value.topics
                         .return_value.create)
        create_method.assert_called_with(body={}, name=EXPANDED_TOPIC)
        create_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value
         .create.return_value.execute.side_effect) = HttpError(
            resp={'status': '409'}, content='')

        try:
            self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC,
                                          fail_if_exists=True)
        except Exception:
            pass  # Expected.
        else:
            self.fail('Topic creation should fail for existing topic when '
                      'fail_if_exists=True')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value
         .get.return_value.execute.side_effect) = HttpError(
            resp={'status': '409'}, content='')

        try:
            self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC,
                                          fail_if_exists=False)
        except Exception:
            self.fail('Topic creation should not fail for existing topic when '
                      'fail_if_exists=False')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_publish(self, mock_service):
        self.pubsub_hook.publish(TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)

        publish_method = (mock_service.return_value.projects.return_value
                          .topics.return_value.publish)
        publish_method.assert_called_with(
            topic=EXPANDED_TOPIC, body={'messages': TEST_MESSAGES})
示例#2
0
def bq_to_pubsub_query_executor(**kwargs):
    """Executes a custom detector query in BigQuery and passes the results to the next task"""

    query = kwargs['templates_dict']['query']
    logging.info(query)
    bigquery_hook = BigQueryHook(use_legacy_sql=False)
    df = bigquery_hook.get_pandas_df(sql=query)

    messages = [{
        'data': b64e(row.to_json().encode()).decode()
    } for index, row in df.iterrows()]
    """splitting the array to 1000 size chunks (PubSub limit)"""
    messages_chunks = chunks(messages, 1000)
    pubsub_hoook = PubSubHook()
    for chunk in messages_chunks:
        pubsub_hoook.publish(project=gcp_project,
                             topic=pubsub_topic,
                             messages=chunk)
示例#3
0
 def execute(self, context):
     hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
                       delegate_to=self.delegate_to)
     hook.publish(self.project, self.topic, self.messages)
示例#4
0
class PubSubHookTest(unittest.TestCase):
    def setUp(self):
        with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
                        new=mock_init):
            self.pubsub_hook = PubSubHook(gcp_conn_id='test')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_topic(self, mock_service):
        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

        create_method = (mock_service.return_value.projects.return_value.
                         topics.return_value.create)
        create_method.assert_called_with(body={}, name=EXPANDED_TOPIC)
        create_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_topic(self, mock_service):
        self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC)

        delete_method = (mock_service.return_value.projects.return_value.
                         topics.return_value.delete)
        delete_method.assert_called_with(topic=EXPANDED_TOPIC)
        delete_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_topic_failifnotexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         delete.return_value.execute.side_effect) = HttpError(
             resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC, True)

        self.assertEqual(str(e.exception),
                         'Topic does not exist: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         create.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC, True)
        self.assertEqual(str(e.exception),
                         'Topic already exists: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         get.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_subscription(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION)

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)
        self.assertEqual(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_different_project_topic(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION, 'a-different-project')

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)

        expected_subscription = 'projects/%s/subscriptions/%s' % (
            'a-different-project', TEST_SUBSCRIPTION)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        create_method.assert_called_with(name=expected_subscription,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)
        self.assertEqual(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_subscription(self, mock_service):
        self.pubsub_hook.delete_subscription(TEST_PROJECT, TEST_SUBSCRIPTION)

        delete_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.delete)
        delete_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION)
        delete_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_subscription_failifnotexists(
            self, mock_service):
        (mock_service.return_value.projects.return_value.subscriptions.
         return_value.delete.return_value.execute.side_effect) = HttpError(
             resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_subscription(TEST_PROJECT,
                                                 TEST_SUBSCRIPTION,
                                                 fail_if_not_exists=True)

        self.assertEqual(
            str(e.exception),
            'Subscription does not exist: %s' % EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    @mock.patch(PUBSUB_STRING.format('uuid4'),
                new_callable=mock.Mock(return_value=lambda: TEST_UUID))
    def test_create_subscription_without_name(self, mock_uuid, mock_service):  # noqa  # pylint: disable=unused-argument,line-too-long
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC)
        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        expected_name = EXPANDED_SUBSCRIPTION.replace(TEST_SUBSCRIPTION,
                                                      'sub-%s' % TEST_UUID)
        create_method.assert_called_with(name=expected_name,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)
        self.assertEqual('sub-%s' % TEST_UUID, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_with_ack_deadline(self, mock_service):
        response = self.pubsub_hook.create_subscription(TEST_PROJECT,
                                                        TEST_TOPIC,
                                                        TEST_SUBSCRIPTION,
                                                        ack_deadline_secs=30)

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 30}
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with(
            num_retries=mock.ANY)
        self.assertEqual(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.subscriptions.
         return_value.create.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_subscription(TEST_PROJECT,
                                                 TEST_TOPIC,
                                                 TEST_SUBSCRIPTION,
                                                 fail_if_exists=True)

        self.assertEqual(
            str(e.exception),
            'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         get.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION)
        self.assertEqual(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_publish(self, mock_service):
        self.pubsub_hook.publish(TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)

        publish_method = (mock_service.return_value.projects.return_value.
                          topics.return_value.publish)
        publish_method.assert_called_with(topic=EXPANDED_TOPIC,
                                          body={'messages': TEST_MESSAGES})

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value.
                       subscriptions.return_value.pull)
        pulled_messages = []
        for i, msg in enumerate(TEST_MESSAGES):
            pulled_messages.append({'ackId': i, 'message': msg})
        pull_method.return_value.execute.return_value = {
            'receivedMessages': pulled_messages
        }

        response = self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
        pull_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION,
                                       body={
                                           'maxMessages': 10,
                                           'returnImmediately': False
                                       })
        self.assertEqual(pulled_messages, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull_no_messages(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value.
                       subscriptions.return_value.pull)
        pull_method.return_value.execute.return_value = {
            'receivedMessages': []
        }

        response = self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
        pull_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION,
                                       body={
                                           'maxMessages': 10,
                                           'returnImmediately': False
                                       })
        self.assertListEqual([], response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull_fails_on_exception(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value.
                       subscriptions.return_value.pull)
        pull_method.return_value.execute.side_effect = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(Exception):
            self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
            pull_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION,
                                           body={
                                               'maxMessages': 10,
                                               'returnImmediately': False
                                           })

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_acknowledge(self, mock_service):
        ack_method = (mock_service.return_value.projects.return_value.
                      subscriptions.return_value.acknowledge)
        self.pubsub_hook.acknowledge(TEST_PROJECT, TEST_SUBSCRIPTION,
                                     ['1', '2', '3'])
        ack_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION,
                                      body={'ackIds': ['1', '2', '3']})

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_acknowledge_fails_on_exception(self, mock_service):
        ack_method = (mock_service.return_value.projects.return_value.
                      subscriptions.return_value.acknowledge)
        ack_method.return_value.execute.side_effect = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(Exception) as e:
            self.pubsub_hook.acknowledge(TEST_PROJECT, TEST_SUBSCRIPTION,
                                         ['1', '2', '3'])
            ack_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION,
                                          body={'ackIds': ['1', '2', '3']})
            print(e)
class PubSubHookTest(unittest.TestCase):
    def setUp(self):
        with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
                        new=mock_init):
            self.pubsub_hook = PubSubHook(gcp_conn_id='test')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_topic(self, mock_service):
        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

        create_method = (mock_service.return_value.projects.return_value.topics
                         .return_value.create)
        create_method.assert_called_with(body={}, name=EXPANDED_TOPIC)
        create_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_topic(self, mock_service):
        self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC)

        delete_method = (mock_service.return_value.projects.return_value.topics
                         .return_value.delete)
        delete_method.assert_called_with(topic=EXPANDED_TOPIC)
        delete_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_topic_failifnotexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics
         .return_value.delete.return_value.execute.side_effect) = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC, True)

        self.assertEquals(str(e.exception),
                          'Topic does not exist: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value
         .create.return_value.execute.side_effect) = HttpError(
            resp={'status': '409'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC, True)
        self.assertEquals(str(e.exception),
                          'Topic already exists: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value
         .get.return_value.execute.side_effect) = HttpError(
            resp={'status': '409'}, content=EMPTY_CONTENT)

        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_subscription(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION)

        create_method = (
            mock_service.return_value.projects.return_value.subscriptions.
            return_value.create)
        expected_body = {
            'topic': EXPANDED_TOPIC,
            'ackDeadlineSeconds': 10
        }
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_different_project_topic(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION, 'a-different-project')

        create_method = (
            mock_service.return_value.projects.return_value.subscriptions.
            return_value.create)

        expected_subscription = 'projects/%s/subscriptions/%s' % (
            'a-different-project', TEST_SUBSCRIPTION)
        expected_body = {
            'topic': EXPANDED_TOPIC,
            'ackDeadlineSeconds': 10
        }
        create_method.assert_called_with(name=expected_subscription,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_subscription(self, mock_service):
        self.pubsub_hook.delete_subscription(TEST_PROJECT, TEST_SUBSCRIPTION)

        delete_method = (mock_service.return_value.projects
                         .return_value.subscriptions.return_value.delete)
        delete_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION)
        delete_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_subscription_failifnotexists(self,
                                                             mock_service):
        (mock_service.return_value.projects.return_value.subscriptions.
         return_value.delete.return_value.execute.side_effect) = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_subscription(
                TEST_PROJECT, TEST_SUBSCRIPTION, fail_if_not_exists=True)

        self.assertEquals(str(e.exception),
                          'Subscription does not exist: %s' %
                          EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    @mock.patch(PUBSUB_STRING.format('uuid4'),
                new_callable=mock.Mock(return_value=lambda: TEST_UUID))
    def test_create_subscription_without_name(self, mock_uuid, mock_service):
        response = self.pubsub_hook.create_subscription(TEST_PROJECT,
                                                        TEST_TOPIC)
        create_method = (
            mock_service.return_value.projects.return_value.subscriptions.
            return_value.create)
        expected_body = {
            'topic': EXPANDED_TOPIC,
            'ackDeadlineSeconds': 10
        }
        expected_name = EXPANDED_SUBSCRIPTION.replace(
            TEST_SUBSCRIPTION, 'sub-%s' % TEST_UUID)
        create_method.assert_called_with(name=expected_name,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals('sub-%s' % TEST_UUID, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_with_ack_deadline(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION, ack_deadline_secs=30)

        create_method = (
            mock_service.return_value.projects.return_value.subscriptions.
            return_value.create)
        expected_body = {
            'topic': EXPANDED_TOPIC,
            'ackDeadlineSeconds': 30
        }
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.
         subscriptions.return_value.create.return_value
         .execute.side_effect) = HttpError(resp={'status': '409'},
                                           content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_subscription(
                TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION,
                fail_if_exists=True)

        self.assertEquals(str(e.exception),
                          'Subscription already exists: %s' %
                          EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value
         .get.return_value.execute.side_effect) = HttpError(
            resp={'status': '409'}, content=EMPTY_CONTENT)

        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION
        )
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_publish(self, mock_service):
        self.pubsub_hook.publish(TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)

        publish_method = (mock_service.return_value.projects.return_value
                          .topics.return_value.publish)
        publish_method.assert_called_with(
            topic=EXPANDED_TOPIC, body={'messages': TEST_MESSAGES})

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value
                       .subscriptions.return_value.pull)
        pulled_messages = []
        for i in range(len(TEST_MESSAGES)):
            pulled_messages.append({'ackId': i, 'message': TEST_MESSAGES[i]})
        pull_method.return_value.execute.return_value = {
            'receivedMessages': pulled_messages}

        response = self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
        pull_method.assert_called_with(
            subscription=EXPANDED_SUBSCRIPTION,
            body={'maxMessages': 10, 'returnImmediately': False})
        self.assertEqual(pulled_messages, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull_no_messages(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value
                       .subscriptions.return_value.pull)
        pull_method.return_value.execute.return_value = {
            'receivedMessages': []}

        response = self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
        pull_method.assert_called_with(
            subscription=EXPANDED_SUBSCRIPTION,
            body={'maxMessages': 10, 'returnImmediately': False})
        self.assertListEqual([], response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_pull_fails_on_exception(self, mock_service):
        pull_method = (mock_service.return_value.projects.return_value
                       .subscriptions.return_value.pull)
        pull_method.return_value.execute.side_effect = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(Exception):
            self.pubsub_hook.pull(TEST_PROJECT, TEST_SUBSCRIPTION, 10)
            pull_method.assert_called_with(
                subscription=EXPANDED_SUBSCRIPTION,
                body={'maxMessages': 10, 'returnImmediately': False})

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_acknowledge(self, mock_service):
        ack_method = (mock_service.return_value.projects.return_value
                      .subscriptions.return_value.acknowledge)
        self.pubsub_hook.acknowledge(
            TEST_PROJECT, TEST_SUBSCRIPTION, ['1', '2', '3'])
        ack_method.assert_called_with(
            subscription=EXPANDED_SUBSCRIPTION,
            body={'ackIds': ['1', '2', '3']})

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_acknowledge_fails_on_exception(self, mock_service):
        ack_method = (mock_service.return_value.projects.return_value
                      .subscriptions.return_value.acknowledge)
        ack_method.return_value.execute.side_effect = HttpError(
            resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(Exception) as e:
            self.pubsub_hook.acknowledge(
                TEST_PROJECT, TEST_SUBSCRIPTION, ['1', '2', '3'])
            ack_method.assert_called_with(
                subscription=EXPANDED_SUBSCRIPTION,
                body={'ackIds': ['1', '2', '3']})
            print(e)
 def execute(self, context):
     hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
                       delegate_to=self.delegate_to)
     hook.publish(self.project, self.topic, self.messages)
示例#7
0
class PubSubHookTest(unittest.TestCase):
    def setUp(self):
        with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
                        new=mock_init):
            self.pubsub_hook = PubSubHook(gcp_conn_id='test')

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_topic(self, mock_service):
        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

        create_method = (mock_service.return_value.projects.return_value.
                         topics.return_value.create)
        create_method.assert_called_with(body={}, name=EXPANDED_TOPIC)
        create_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_topic(self, mock_service):
        self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC)

        delete_method = (mock_service.return_value.projects.return_value.
                         topics.return_value.delete)
        delete_method.assert_called_with(topic=EXPANDED_TOPIC)
        delete_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_topic_failifnotexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         delete.return_value.execute.side_effect) = HttpError(
             resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_topic(TEST_PROJECT, TEST_TOPIC, True)

        self.assertEquals(str(e.exception),
                          'Topic does not exist: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         create.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC, True)
        self.assertEquals(str(e.exception),
                          'Topic already exists: %s' % EXPANDED_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_preexisting_topic_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         get.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        self.pubsub_hook.create_topic(TEST_PROJECT, TEST_TOPIC)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_nonexistent_subscription(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION)

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_different_project_topic(self, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION, 'a-different-project')

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)

        expected_subscription = 'projects/%s/subscriptions/%s' % (
            'a-different-project', TEST_SUBSCRIPTION)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        create_method.assert_called_with(name=expected_subscription,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_subscription(self, mock_service):
        self.pubsub_hook.delete_subscription(TEST_PROJECT, TEST_SUBSCRIPTION)

        delete_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.delete)
        delete_method.assert_called_with(subscription=EXPANDED_SUBSCRIPTION)
        delete_method.return_value.execute.assert_called_with()

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_delete_nonexisting_subscription_failifnotexists(
            self, mock_service):
        (mock_service.return_value.projects.return_value.subscriptions.
         return_value.delete.return_value.execute.side_effect) = HttpError(
             resp={'status': '404'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.delete_subscription(TEST_PROJECT,
                                                 TEST_SUBSCRIPTION,
                                                 fail_if_not_exists=True)

        self.assertEquals(
            str(e.exception),
            'Subscription does not exist: %s' % EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    @mock.patch(PUBSUB_STRING.format('uuid4'),
                new_callable=mock.Mock(return_value=lambda: TEST_UUID))
    def test_create_subscription_without_name(self, mock_uuid, mock_service):
        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC)
        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 10}
        expected_name = EXPANDED_SUBSCRIPTION.replace(TEST_SUBSCRIPTION,
                                                      'sub-%s' % TEST_UUID)
        create_method.assert_called_with(name=expected_name,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals('sub-%s' % TEST_UUID, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_with_ack_deadline(self, mock_service):
        response = self.pubsub_hook.create_subscription(TEST_PROJECT,
                                                        TEST_TOPIC,
                                                        TEST_SUBSCRIPTION,
                                                        ack_deadline_secs=30)

        create_method = (mock_service.return_value.projects.return_value.
                         subscriptions.return_value.create)
        expected_body = {'topic': EXPANDED_TOPIC, 'ackDeadlineSeconds': 30}
        create_method.assert_called_with(name=EXPANDED_SUBSCRIPTION,
                                         body=expected_body)
        create_method.return_value.execute.assert_called_with()
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_failifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.subscriptions.
         return_value.create.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        with self.assertRaises(PubSubException) as e:
            self.pubsub_hook.create_subscription(TEST_PROJECT,
                                                 TEST_TOPIC,
                                                 TEST_SUBSCRIPTION,
                                                 fail_if_exists=True)

        self.assertEquals(
            str(e.exception),
            'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_create_subscription_nofailifexists(self, mock_service):
        (mock_service.return_value.projects.return_value.topics.return_value.
         get.return_value.execute.side_effect) = HttpError(
             resp={'status': '409'}, content=EMPTY_CONTENT)

        response = self.pubsub_hook.create_subscription(
            TEST_PROJECT, TEST_TOPIC, TEST_SUBSCRIPTION)
        self.assertEquals(TEST_SUBSCRIPTION, response)

    @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
    def test_publish(self, mock_service):
        self.pubsub_hook.publish(TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)

        publish_method = (mock_service.return_value.projects.return_value.
                          topics.return_value.publish)
        publish_method.assert_called_with(topic=EXPANDED_TOPIC,
                                          body={'messages': TEST_MESSAGES})