def test_creation(self):
# Test the case where queue needs to create the topic.
pubsub = Mock()
topic = Mock()
topic.exists.return_value = False
pubsub.topic.return_value = topic
q = Queue(pubsub)
assert pubsub.topic.called
assert topic.exists.called
assert q.topic == topic
sub = Mock()
sub.exists.return_value = False
# Test the case where it needs to create the subcription.
with patch('gcloud.pubsub.Subscription') as SubscriptionMock:
SubscriptionMock.return_value = sub
rsub = q._get_or_create_subscription()
assert rsub == sub
assert SubscriptionMock.called_with('psq-default-shared', topic)
assert sub.exists.called
assert sub.create.called
# Test case where subscription exists and it should re-use it.
with patch('gcloud.pubsub.Subscription') as SubscriptionMock:
sub.reset_mock()
SubscriptionMock.return_value = sub
sub.exists.return_value = True
rsub = q._get_or_create_subscription()
assert rsub == sub
assert not sub.create.called
def test_constructor_creates_topic():
publisher_client = make_publisher_client()
subscriber_client = make_subscriber_client()
publisher_client.get_topic.side_effect = (google.cloud.exceptions.NotFound(
None, None))
q = Queue(publisher_client, subscriber_client, 'test-project')
publisher_client.create_topic.assert_called_once_with(q._get_topic_path())
def test__pubsub_message_callback_bad_value():
callback = mock.Mock()
message = mock.create_autospec(
google.cloud.pubsub_v1.subscriber.message.Message, instance=True)
message.data = b'bad'
Queue._pubsub_message_callback(callback, message)
assert not callback.called
assert message.ack.called
def test_queue(self):
# Test queueing tasks.
with patch('psq.queue.Queue._get_or_create_topic'):
q = Queue(Mock())
q.storage.put_task = Mock()
r = q.enqueue(sum, 1, 2, arg='c')
assert q.topic.publish.called
assert q.storage.put_task.called
t = q.storage.put_task.call_args[0][0]
assert t.f == sum
assert t.args == (1, 2)
assert t.kwargs == {'arg': 'c'}
assert r.task_id == t.id
def test_queue():
# Test queueing tasks.
with patch('psq.queue.Queue._get_or_create_topic'):
q = Queue(Mock())
q.storage.put_task = Mock()
r = q.enqueue(sum, 1, 2, arg='c')
assert q.topic.publish.called
assert q.storage.put_task.called
t = q.storage.put_task.call_args[0][0]
assert t.f == sum
assert t.args == (1, 2)
assert t.kwargs == {'arg': 'c'}
assert r.task_id == t.id
def test_constructor_existing_topic():
publisher_client = make_publisher_client()
subscriber_client = make_subscriber_client()
Queue(publisher_client, subscriber_client, 'test-project')
publisher_client.create_topic.assert_not_called()
def test_get_or_create_subscription():
pubsub = Mock()
topic = Mock()
topic.exists.return_value = True
pubsub.topic.return_value = topic
q = Queue(pubsub)
sub = Mock()
sub.exists.return_value = False
# Test the case where it needs to create the subcription.
with patch('google.cloud.pubsub.Subscription') as SubscriptionMock:
SubscriptionMock.return_value = sub
rsub = q._get_or_create_subscription()
assert rsub == sub
assert SubscriptionMock.called_with('psq-default-shared', topic)
assert sub.exists.called
assert sub.create.called
# Test case where subscription exists and it should re-use it.
with patch('google.cloud.pubsub.Subscription') as SubscriptionMock:
sub.reset_mock()
SubscriptionMock.return_value = sub
sub.exists.return_value = True
rsub = q._get_or_create_subscription()
assert rsub == sub
assert not sub.create.called
# Test case where subscription gets created after we check that it
# doesn't exist.
with patch('google.cloud.pubsub.Subscription') as SubscriptionMock:
sub.reset_mock()
SubscriptionMock.return_value = sub
sub.exists.return_value = False
sub.create.side_effect = google.cloud.exceptions.Conflict('')
rsub = q._get_or_create_subscription()
assert sub.create.called
assert rsub == sub
def test_get_or_create_subscription():
pubsub = Mock()
topic = Mock()
topic.exists.return_value = True
pubsub.topic.return_value = topic
q = Queue(pubsub)
sub = Mock()
sub.exists.return_value = False
# Test the case where it needs to create the subcription.
with patch("google.cloud.pubsub.Subscription") as SubscriptionMock:
SubscriptionMock.return_value = sub
rsub = q._get_or_create_subscription()
assert rsub == sub
assert SubscriptionMock.called_with("psq-default-shared", topic)
assert sub.exists.called
assert sub.create.called
# Test case where subscription exists and it should re-use it.
with patch("google.cloud.pubsub.Subscription") as SubscriptionMock:
sub.reset_mock()
SubscriptionMock.return_value = sub
sub.exists.return_value = True
rsub = q._get_or_create_subscription()
assert rsub == sub
assert not sub.create.called
# Test case where subscription gets created after we check that it
# doesn't exist.
with patch("google.cloud.pubsub.Subscription") as SubscriptionMock:
sub.reset_mock()
SubscriptionMock.return_value = sub
sub.exists.return_value = False
sub.create.side_effect = google.cloud.exceptions.Conflict("")
rsub = q._get_or_create_subscription()
assert sub.create.called
assert rsub == sub
def test_dequeue(self):
# Test dequeueing (fetching) tasks.
with patch('psq.queue.Queue._get_or_create_topic'):
q = Queue(Mock())
t = Task('1', sum, (1, 2), {'arg': 'c'})
sub_mock = Mock()
q._get_or_create_subscription = Mock(return_value=sub_mock)
# No messages
sub_mock.pull.return_value = []
tasks = q.dequeue()
assert sub_mock.pull.called
assert not tasks
# One Message
sub_mock.pull.reset_mock()
sub_mock.pull.return_value = [
('ack_id', MockMessage(dumps(t)))]
tasks = q.dequeue()
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(['ack_id'])
assert tasks[0].id == t.id
assert tasks[0].f == t.f
assert tasks[0].args == t.args
assert tasks[0].kwargs == t.kwargs
# Bad message
sub_mock.pull.reset_mock()
sub_mock.acknowledge.reset_mock()
sub_mock.pull.return_value = [
('ack_id', MockMessage('this is a bad pickle string'))]
tasks = q.dequeue()
assert not tasks
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(['ack_id'])
def test_context():
# Test queue-local context.
pubsub = Mock()
pubsub.topic.return_value = Mock()
q = Queue(pubsub)
with q.queue_context():
assert current_queue == q
# Test additional context manager.
spy = Mock()
@contextmanager
def extra_context():
spy()
yield
q.extra_context = extra_context
with q.queue_context():
assert spy.called
def test_context(self):
# Test queue-local context.
pubsub = Mock()
pubsub.topic.return_value = Mock()
q = Queue(pubsub)
with q.queue_context():
assert current_queue == q
# Test additional context manager.
spy = Mock()
@contextmanager
def extra_context():
spy()
yield
q.extra_context = extra_context
with q.queue_context():
assert spy.called
def test_creation_existing_topic():
# Test the case where queue needs to create the topic.
pubsub = Mock()
topic = Mock()
topic.exists.return_value = False
topic.create.side_effect = google.cloud.exceptions.Conflict('')
pubsub.topic.return_value = topic
q = Queue(pubsub)
assert pubsub.topic.called
assert topic.create.called
assert q.topic == topic
def test_dequeue():
# Test dequeueing (fetching) tasks.
with patch("psq.queue.Queue._get_or_create_topic"):
q = Queue(Mock())
t = Task("1", sum, (1, 2), {"arg": "c"})
sub_mock = Mock()
q._get_or_create_subscription = Mock(return_value=sub_mock)
# No messages
sub_mock.pull.return_value = []
tasks = q.dequeue()
assert sub_mock.pull.called
assert not tasks
# One Message
sub_mock.pull.reset_mock()
sub_mock.pull.return_value = [("ack_id", MockMessage(dumps(t)))]
tasks = q.dequeue()
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(["ack_id"])
assert tasks[0].id == t.id
assert tasks[0].f == t.f
assert tasks[0].args == t.args
assert tasks[0].kwargs == t.kwargs
# Bad message
sub_mock.pull.reset_mock()
sub_mock.acknowledge.reset_mock()
sub_mock.pull.return_value = [("ack_id", MockMessage("this is a bad pickle string"))]
tasks = q.dequeue()
assert not tasks
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(["ack_id"])
def test_creation():
# Test the case where queue needs to create the topic.
pubsub = Mock()
topic = Mock()
topic.exists.return_value = False
pubsub.topic.return_value = topic
q = Queue(pubsub)
assert pubsub.topic.called
assert topic.create.called
assert q.topic == topic
sub = Mock()
sub.exists.return_value = False
def test_dequeue():
# Test dequeueing (fetching) tasks.
with patch('psq.queue.Queue._get_or_create_topic'):
q = Queue(Mock())
t = Task('1', sum, (1, 2), {'arg': 'c'})
sub_mock = Mock()
q._get_or_create_subscription = Mock(return_value=sub_mock)
# No messages
sub_mock.pull.return_value = []
tasks = q.dequeue()
assert sub_mock.pull.called
assert not tasks
# One Message
sub_mock.pull.reset_mock()
sub_mock.pull.return_value = [('ack_id', MockMessage(dumps(t)))]
tasks = q.dequeue()
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(['ack_id'])
assert tasks[0].id == t.id
assert tasks[0].f == t.f
assert tasks[0].args == t.args
assert tasks[0].kwargs == t.kwargs
# Bad message
sub_mock.pull.reset_mock()
sub_mock.acknowledge.reset_mock()
sub_mock.pull.return_value = [
('ack_id', MockMessage('this is a bad pickle string'))
]
tasks = q.dequeue()
assert not tasks
assert sub_mock.pull.called
sub_mock.acknowledge.assert_called_once_with(['ack_id'])
def test_string_function():
q = Queue(pubsub=None, storage=TestStorage(), async=False)
r = q.enqueue('psq.queue_test.dummy_queue_func')
assert r.result() == "Hello"
def test_synchronous_fail():
q = Queue(pubsub=None, storage=TestStorage(), async=False)
r = q.enqueue(sum, "2")
with pytest.raises(TypeError):
r.result()
def test_synchronous_success():
q = Queue(pubsub=None, storage=TestStorage(), async=False)
r = q.enqueue(sum, [1, 2])
assert r.result() == 3
def test_cleanup():
pubsub = Mock()
pubsub.topic.return_value = Mock()
q = Queue(pubsub)
q.cleanup()
def make_queue(**kwargs):
publisher_client = make_publisher_client()
subscriber_client = make_subscriber_client()
return Queue(publisher_client, subscriber_client, 'test-project', **kwargs)
def test_synchronous_success():
q = Queue(pubsub=None, storage=TestStorage(), async=False)
r = q.enqueue(sum, [1, 2])
assert r.result() == 3
def test_cleanup(self):
pubsub = Mock()
pubsub.topic.return_value = Mock()
q = Queue(pubsub)
q.cleanup()
def test_synchronous_fail():
q = Queue(pubsub=None, storage=TestStorage(), async=False)
r = q.enqueue(sum, "2")
with pytest.raises(TypeError):
r.result()
