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_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():
# 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_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_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"])
