Consumer: __exit__ should not cancel if connection error (Closes #670)

This commit is contained in:
Ask Solem 2016-12-07 14:28:16 -08:00
parent e03cbbe0d8
commit 8a7ea091bf
2 changed files with 17 additions and 5 deletions

View File

@ -430,11 +430,14 @@ class Consumer(object):
self.consume()
return self
def __exit__(self, *exc_info):
try:
self.cancel()
except Exception:
pass
def __exit__(self, exc_type, exc_val, exc_tb):
if self.channel:
conn_errors = self.channel.connection.client.connection_errors
if not isinstance(exc_val, conn_errors):
try:
self.cancel()
except Exception:
pass
def add_queue(self, queue):
"""Add a queue to the list of queues to consume from.

View File

@ -262,6 +262,15 @@ class test_Consumer:
pass
c.cancel.assert_called_with()
def test_enter_exit_cancel_not_called_on_connection_error(self):
c = Consumer(self.connection)
c.cancel = Mock(name='Consumer.cancel')
assert self.connection.connection_errors
with pytest.raises(self.connection.connection_errors[0]):
with c:
raise self.connection.connection_errors[0]()
c.cancel.assert_not_called()
def test_receive_callback_accept(self):
message = Mock(name='Message')
message.errors = []