def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
hook.delete_subscription(self.project,
self.subscription,
fail_if_not_exists=self.fail_if_not_exists)
def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
hook.create_topic(self.project,
self.topic,
fail_if_exists=self.fail_if_exists)
def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
return hook.create_subscription(
self.topic_project, self.topic, self.subscription,
self.subscription_project, self.ack_deadline_secs,
self.fail_if_exists)
def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
return hook.create_subscription(
self.topic_project, self.topic, self.subscription,
self.subscription_project, self.ack_deadline_secs,
self.fail_if_exists)
def poke(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
self._messages = hook.pull(self.project, self.subscription,
self.max_messages, self.return_immediately)
if self._messages and self.ack_messages:
ack_ids = [m['ackId'] for m in self._messages if m.get('ackId')]
hook.acknowledge(self.project, self.subscription, ack_ids)
return self._messages
def poke(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
self._messages = hook.pull(
self.project, self.subscription, self.max_messages,
self.return_immediately)
if self._messages and self.ack_messages:
if self.ack_messages:
ack_ids = [m['ackId'] for m in self._messages if m.get('ackId')]
hook.acknowledge(self.project, self.subscription, ack_ids)
return self._messages
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)
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})
def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
hook.create_topic(self.project, self.topic,
fail_if_exists=self.fail_if_exists)
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)
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.pubsub_hook = PubSubHook(gcp_conn_id='test')
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)
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.pubsub_hook = PubSubHook(gcp_conn_id='test')
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)
def execute(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
hook.delete_subscription(self.project, self.subscription,
fail_if_not_exists=self.fail_if_not_exists)
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})
