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,7 +430,10 @@ class Consumer(object):
self.consume() self.consume()
return self return self
def __exit__(self, *exc_info): 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: try:
self.cancel() self.cancel()
except Exception: except Exception:

View File

@ -262,6 +262,15 @@ class test_Consumer:
pass pass
c.cancel.assert_called_with() 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): def test_receive_callback_accept(self):
message = Mock(name='Message') message = Mock(name='Message')
message.errors = [] message.errors = []