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 poke(self, context):
hook = PubSubHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to)
pulled_messages = hook.pull(project_id=self.project_id,
subscription=self.subscription,
max_messages=self.max_messages,
return_immediately=self.return_immediately)
self._messages = [MessageToDict(m) for m in pulled_messages]
if self._messages and self.ack_messages:
ack_ids = [m['ackId'] for m in self._messages if m.get('ackId')]
hook.acknowledge(project_id=self.project_id,
subscription=self.subscription,
ack_ids=ack_ids)
return self._messages
class TestPubSubHook(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("airflow.gcp.hooks.pubsub.PubSubHook.client_info", new_callable=mock.PropertyMock)
@mock.patch("airflow.gcp.hooks.pubsub.PubSubHook._get_credentials")
@mock.patch("airflow.gcp.hooks.pubsub.PublisherClient")
def test_publisher_client_creation(self, mock_client, mock_get_creds, mock_client_info):
self.assertIsNone(self.pubsub_hook._client)
result = self.pubsub_hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value
)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.pubsub_hook._client, result)
@mock.patch("airflow.gcp.hooks.pubsub.PubSubHook.client_info", new_callable=mock.PropertyMock)
@mock.patch("airflow.gcp.hooks.pubsub.PubSubHook._get_credentials")
@mock.patch("airflow.gcp.hooks.pubsub.SubscriberClient")
def test_subscriber_client_creation(self, mock_client, mock_get_creds, mock_client_info):
self.assertIsNone(self.pubsub_hook._client)
result = self.pubsub_hook.subscriber_client
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value
)
self.assertEqual(mock_client.return_value, result)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_create_nonexistent_topic(self, mock_service):
create_method = mock_service.return_value.create_topic
self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC)
create_method.assert_called_once_with(
name=EXPANDED_TOPIC,
labels=LABELS,
message_storage_policy=None,
kms_key_name=None,
retry=None,
timeout=None,
metadata=None
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_delete_topic(self, mock_service):
delete_method = mock_service.return_value.delete_topic
self.pubsub_hook.delete_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC)
delete_method.assert_called_once_with(
topic=EXPANDED_TOPIC,
retry=None,
timeout=None,
metadata=None
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_delete_nonexisting_topic_failifnotexists(self, mock_service):
mock_service.return_value.delete_topic.side_effect = NotFound(
'Topic does not exists: %s' % EXPANDED_TOPIC
)
with self.assertRaises(PubSubException) as e:
self.pubsub_hook.delete_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_not_exists=True)
self.assertEqual(str(e.exception), 'Topic does not exist: %s' % EXPANDED_TOPIC)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_delete_topic_api_call_error(self, mock_service):
mock_service.return_value.delete_topic.side_effect = GoogleAPICallError(
'Error deleting topic: %s' % EXPANDED_TOPIC
)
with self.assertRaises(PubSubException):
self.pubsub_hook.delete_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_not_exists=True)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_create_preexisting_topic_failifexists(self, mock_service):
mock_service.return_value.create_topic.side_effect = AlreadyExists(
'Topic already exists: %s' % TEST_TOPIC
)
with self.assertRaises(PubSubException) as e:
self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_exists=True)
self.assertEqual(str(e.exception), 'Topic already exists: %s' % TEST_TOPIC)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_create_preexisting_topic_nofailifexists(self, mock_service):
mock_service.return_value.create_topic.side_effect = AlreadyExists(
'Topic already exists: %s' % EXPANDED_TOPIC
)
self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_create_topic_api_call_error(self, mock_service):
mock_service.return_value.create_topic.side_effect = GoogleAPICallError(
'Error creating topic: %s' % TEST_TOPIC
)
with self.assertRaises(PubSubException):
self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_exists=True)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_nonexistent_subscription(self, mock_service):
create_method = mock_service.create_subscription
response = self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION
)
create_method.assert_called_once_with(
name=EXPANDED_SUBSCRIPTION,
topic=EXPANDED_TOPIC,
push_config=None,
ack_deadline_seconds=10,
retain_acked_messages=None,
message_retention_duration=None,
labels=LABELS,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(TEST_SUBSCRIPTION, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_subscription_different_project_topic(self, mock_service):
create_method = mock_service.create_subscription
response = self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT,
topic=TEST_TOPIC,
subscription=TEST_SUBSCRIPTION,
subscription_project_id='a-different-project'
)
expected_subscription = 'projects/{}/subscriptions/{}'.format(
'a-different-project', TEST_SUBSCRIPTION
)
create_method.assert_called_once_with(
name=expected_subscription,
topic=EXPANDED_TOPIC,
push_config=None,
ack_deadline_seconds=10,
retain_acked_messages=None,
message_retention_duration=None,
labels=LABELS,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(TEST_SUBSCRIPTION, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_delete_subscription(self, mock_service):
self.pubsub_hook.delete_subscription(project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION)
delete_method = mock_service.delete_subscription
delete_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
retry=None,
timeout=None,
metadata=None
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_delete_nonexisting_subscription_failifnotexists(self, mock_service):
mock_service.delete_subscription.side_effect = NotFound(
'Subscription does not exists: %s' % EXPANDED_SUBSCRIPTION
)
with self.assertRaises(PubSubException) as e:
self.pubsub_hook.delete_subscription(
project_id=TEST_PROJECT, subscription=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.subscriber_client'))
def test_delete_subscription_api_call_error(self, mock_service):
mock_service.delete_subscription.side_effect = GoogleAPICallError(
'Error deleting subscription %s' % EXPANDED_SUBSCRIPTION
)
with self.assertRaises(PubSubException):
self.pubsub_hook.delete_subscription(
project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, fail_if_not_exists=True
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
@mock.patch(PUBSUB_STRING.format('uuid4'), new_callable=mock.Mock(return_value=lambda: TEST_UUID))
def test_create_subscription_without_subscription_name(self, mock_uuid, mock_service): # noqa # pylint: disable=unused-argument,line-too-long
create_method = mock_service.create_subscription
expected_name = EXPANDED_SUBSCRIPTION.replace(TEST_SUBSCRIPTION, 'sub-%s' % TEST_UUID)
response = self.pubsub_hook.create_subscription(project_id=TEST_PROJECT, topic=TEST_TOPIC)
create_method.assert_called_once_with(
name=expected_name,
topic=EXPANDED_TOPIC,
push_config=None,
ack_deadline_seconds=10,
retain_acked_messages=None,
message_retention_duration=None,
labels=LABELS,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual('sub-%s' % TEST_UUID, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_subscription_with_ack_deadline(self, mock_service):
create_method = mock_service.create_subscription
response = self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION, ack_deadline_secs=30
)
create_method.assert_called_once_with(
name=EXPANDED_SUBSCRIPTION,
topic=EXPANDED_TOPIC,
push_config=None,
ack_deadline_seconds=30,
retain_acked_messages=None,
message_retention_duration=None,
labels=LABELS,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(TEST_SUBSCRIPTION, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_subscription_failifexists(self, mock_service):
mock_service.create_subscription.side_effect = AlreadyExists(
'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION
)
with self.assertRaises(PubSubException) as e:
self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION, fail_if_exists=True
)
self.assertEqual(str(e.exception), 'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_subscription_api_call_error(self, mock_service):
mock_service.create_subscription.side_effect = GoogleAPICallError(
'Error creating subscription %s' % EXPANDED_SUBSCRIPTION
)
with self.assertRaises(PubSubException):
self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION, fail_if_exists=True
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_create_subscription_nofailifexists(self, mock_service):
mock_service.create_subscription.side_effect = AlreadyExists(
'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION
)
response = self.pubsub_hook.create_subscription(
project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION
)
self.assertEqual(TEST_SUBSCRIPTION, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_publish(self, mock_service):
publish_method = mock_service.return_value.publish
self.pubsub_hook.publish(project_id=TEST_PROJECT, topic=TEST_TOPIC, messages=TEST_MESSAGES)
calls = [
mock.call(topic=EXPANDED_TOPIC, data=message.get("data", b''), **message.get('attributes', {}))
for message in TEST_MESSAGES
]
publish_method.has_calls(calls)
@mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn'))
def test_publish_api_call_error(self, mock_service):
publish_method = mock_service.return_value.publish
publish_method.side_effect = GoogleAPICallError(
'Error publishing to topic {}'.format(EXPANDED_SUBSCRIPTION)
)
with self.assertRaises(PubSubException):
self.pubsub_hook.publish(project_id=TEST_PROJECT, topic=TEST_TOPIC, messages=TEST_MESSAGES)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_pull(self, mock_service):
pull_method = mock_service.pull
pulled_messages = []
for i, msg in enumerate(TEST_MESSAGES):
pulled_messages.append({'ackId': i, 'message': msg})
pull_method.return_value.received_messages = pulled_messages
response = self.pubsub_hook.pull(
project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=10
)
pull_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
max_messages=10,
return_immediately=False,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(pulled_messages, response)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_pull_no_messages(self, mock_service):
pull_method = mock_service.pull
pull_method.return_value.received_messages = []
response = self.pubsub_hook.pull(
project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=10
)
pull_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
max_messages=10,
return_immediately=False,
retry=None,
timeout=None,
metadata=None,
)
self.assertListEqual([], response)
@parameterized.expand([
(exception, ) for exception in [
HttpError(resp={'status': '404'}, content=EMPTY_CONTENT),
GoogleAPICallError("API Call Error")
]
])
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_pull_fails_on_exception(self, exception, mock_service):
pull_method = mock_service.pull
pull_method.side_effect = exception
with self.assertRaises(PubSubException):
self.pubsub_hook.pull(project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=10)
pull_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
max_messages=10,
return_immediately=False,
retry=None,
timeout=None,
metadata=None,
)
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_acknowledge(self, mock_service):
ack_method = mock_service.acknowledge
self.pubsub_hook.acknowledge(
project_id=TEST_PROJECT,
subscription=TEST_SUBSCRIPTION,
ack_ids=['1', '2', '3']
)
ack_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
ack_ids=['1', '2', '3'],
retry=None,
timeout=None,
metadata=None
)
@parameterized.expand([
(exception, ) for exception in [
HttpError(resp={'status': '404'}, content=EMPTY_CONTENT),
GoogleAPICallError("API Call Error")
]
])
@mock.patch(PUBSUB_STRING.format('PubSubHook.subscriber_client'))
def test_acknowledge_fails_on_exception(self, exception, mock_service):
ack_method = mock_service.acknowledge
ack_method.side_effect = exception
with self.assertRaises(PubSubException):
self.pubsub_hook.acknowledge(
project_id=TEST_PROJECT,
subscription=TEST_SUBSCRIPTION,
ack_ids=['1', '2', '3']
)
ack_method.assert_called_once_with(
subscription=EXPANDED_SUBSCRIPTION,
ack_ids=['1', '2', '3'],
retry=None,
timeout=None,
metadata=None
)
@parameterized.expand([
(messages, ) for messages in [
[{"data": b'test'}],
[{"data": b''}],
[{"data": b'test', "attributes": {"weight": "100kg"}}],
[{"data": b'', "attributes": {"weight": "100kg"}}],
[{"attributes": {"weight": "100kg"}}],
]
])
def test_messages_validation_positive(self, messages):
PubSubHook._validate_messages(messages)
@parameterized.expand([
([("wrong type", )], "Wrong message type. Must be a dictionary."),
([{"wrong_key": b'test'}], "Wrong message. Dictionary must contain 'data' or 'attributes'."),
([{"data": 'wrong string'}], "Wrong message. 'data' must be send as a bytestring"),
([{"data": None}], "Wrong message. 'data' must be send as a bytestring"),
(
[{"attributes": None}],
"Wrong message. If 'data' is not provided 'attributes' must be a non empty dictionary."
),
(
[{"attributes": "wrong string"}],
"Wrong message. If 'data' is not provided 'attributes' must be a non empty dictionary."
)
])
def test_messages_validation_negative(self, messages, error_message):
with self.assertRaises(PubSubException) as e:
PubSubHook._validate_messages(messages)
self.assertEqual(str(e.exception), error_message)
class TestPubSubHook(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_once_with(body={}, name=EXPANDED_TOPIC)
create_method.return_value.execute.assert_called_once_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_once_with(topic=EXPANDED_TOPIC)
delete_method.return_value.execute.assert_called_once_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_once_with(name=EXPANDED_SUBSCRIPTION,
body=expected_body)
create_method.return_value.execute.assert_called_once_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_once_with(name=expected_subscription,
body=expected_body)
create_method.return_value.execute.assert_called_once_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_once_with(
subscription=EXPANDED_SUBSCRIPTION)
delete_method.return_value.execute.assert_called_once_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_once_with(name=expected_name,
body=expected_body)
create_method.return_value.execute.assert_called_once_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_once_with(name=EXPANDED_SUBSCRIPTION,
body=expected_body)
create_method.return_value.execute.assert_called_once_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_once_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_once_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_once_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_once_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_once_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_once_with(
subscription=EXPANDED_SUBSCRIPTION,
body={'ackIds': ['1', '2', '3']})
print(e)