User initiated close (#150)

Author: eandersson

amqpstorm/channel.py

@@ -69,6 +69,7 @@ def __init__(
         self._exchange = Exchange(self)
         self._tx = Tx(self)
         self._queue = Queue(self)
+        self._user_closed: bool = False
 
     def __enter__(self) -> Channel:
         return self
@@ -173,7 +174,12 @@ def build_inbound_messages(
 
         :rtype: :py:class:`generator`
         """
-        self.check_for_errors()
+        try:
+            self.check_for_errors()
+        except (AMQPConnectionError, AMQPChannelError):
+            if self._user_initiated_close():
+                return
+            raise
         if message_impl:
             if not issubclass(message_impl, BaseMessage):
                 raise AMQPInvalidArgument(
@@ -182,10 +188,20 @@ def build_inbound_messages(
         else:
             message_impl = Message
         while not self.is_closed:
-            message = self._build_message(auto_decode=auto_decode,
-                                          message_impl=message_impl)
+            try:
+                message = self._build_message(auto_decode=auto_decode,
+                                              message_impl=message_impl)
+            except (AMQPConnectionError, AMQPChannelError):
+                if self._user_initiated_close():
+                    return
+                raise
             if not message:
-                self.check_for_errors()
+                try:
+                    self.check_for_errors()
+                except (AMQPConnectionError, AMQPChannelError):
+                    if self._user_initiated_close():
+                        return
+                    raise
                 time.sleep(IDLE_WAIT)
                 if break_on_empty and not self._inbound:
                     break
@@ -195,6 +211,9 @@ def build_inbound_messages(
                 continue
             yield message
 
+    def _user_initiated_close(self) -> bool:
+        return self._user_closed or self._connection._user_closed
+
     def close(self, reply_code: int = 200, reply_text: str = '') -> None:
         """Close Channel.
 
@@ -212,6 +231,7 @@ def close(self, reply_code: int = 200, reply_text: str = '') -> None:
             raise AMQPInvalidArgument('reply_code should be an integer')
         elif not compatibility.is_string(reply_text):
             raise AMQPInvalidArgument('reply_text should be a string')
+        self._user_closed = True
         try:
             if self._connection.is_closed or not self.is_open:
                 self.stop_consuming()

amqpstorm/connection.py

@@ -110,6 +110,7 @@ def __init__(
         self._channel0 = Channel0(self, self.parameters['client_properties'])
         self._channels: dict[int, Channel] = {}
         self._last_channel_id: int | None = None
+        self._user_closed: bool = False
         self.heartbeat = Heartbeat(self.parameters['heartbeat'],
                                    self._channel0.send_heartbeat)
         if not kwargs.get('lazy', False):
@@ -240,6 +241,7 @@ def close(self) -> None:
         :return:
         """
         LOGGER.debug('Connection Closing')
+        self._user_closed = True
         if not self.is_closed:
             self.set_state(self.CLOSING)
         self.heartbeat.stop()

amqpstorm/tests/unit/channel/test_channel.py

@@ -79,6 +79,20 @@ def on_close_ok(_, frame_out):
         self.assertEqual(channel._state, channel.CLOSED)
         self.assertFalse(channel.exceptions)
 
+    def test_channel_close_sets_user_closed_flag(self):
+        def on_close_ok(_, frame_out):
+            if isinstance(frame_out, commands.Basic.Cancel):
+                channel.rpc.on_frame(commands.Basic.CancelOk())
+                return
+            channel.rpc.on_frame(commands.Channel.CloseOk())
+
+        channel = Channel(0, FakeConnection(on_write=on_close_ok), 360)
+        channel.set_state(channel.OPEN)
+
+        self.assertFalse(channel._user_closed)
+        channel.close()
+        self.assertTrue(channel._user_closed)
+
     def test_channel_close_gracefully_with_queued_error(self):
         def on_close_ok(_, frame_out):
             if isinstance(frame_out, commands.Basic.Cancel):

amqpstorm/tests/unit/channel/test_channel_message_handling.py

@@ -8,6 +8,7 @@
 from pamqp.body import ContentBody
 
 from amqpstorm import AMQPChannelError
+from amqpstorm import AMQPConnectionError
 from amqpstorm import Channel
 from amqpstorm import Message
 from amqpstorm.tests.utility import FakeConnection
@@ -338,6 +339,105 @@ def test_channel_build_inbound_messages_as_tuple(self):
 
         self.assertEqual(messages_consumed, 1)
 
+    def test_channel_build_inbound_messages_returns_on_user_close(self):
+        channel = Channel(0, FakeConnection(), 360)
+        channel.set_state(channel.OPEN)
+
+        message = self.message.encode('utf-8')
+        deliver = commands.Basic.Deliver()
+        header = ContentHeader(body_size=len(message))
+        body = ContentBody(value=message)
+        channel._inbound = collections.deque([deliver, header, body])
+
+        consumed = 0
+        for _ in channel.build_inbound_messages(break_on_empty=False):
+            consumed += 1
+            channel._user_closed = True
+            channel.set_state(channel.CLOSED)
+        self.assertEqual(consumed, 1)
+        self.assertFalse(channel.exceptions)
+
+    def test_channel_build_inbound_messages_returns_on_connection_user_close(self):
+        connection = FakeConnection()
+        channel = Channel(0, connection, 360)
+        channel.set_state(channel.OPEN)
+
+        message = self.message.encode('utf-8')
+        deliver = commands.Basic.Deliver()
+        header = ContentHeader(body_size=len(message))
+        body = ContentBody(value=message)
+        channel._inbound = collections.deque([deliver, header, body])
+
+        consumed = 0
+        for _ in channel.build_inbound_messages(break_on_empty=False):
+            consumed += 1
+            connection._user_closed = True
+            connection.set_state(connection.CLOSED)
+        self.assertEqual(consumed, 1)
+
+    def test_channel_build_inbound_messages_raises_on_server_close(self):
+        channel = Channel(0, FakeConnection(), 360)
+        channel.set_state(channel.OPEN)
+
+        message = self.message.encode('utf-8')
+        deliver = commands.Basic.Deliver()
+        header = ContentHeader(body_size=len(message))
+        body = ContentBody(value=message)
+        channel._inbound = collections.deque([deliver, header, body])
+
+        def drive():
+            for _ in channel.build_inbound_messages(break_on_empty=False):
+                channel.exceptions.append(
+                    AMQPChannelError('closed by server', reply_code=406)
+                )
+
+        self.assertRaises(AMQPChannelError, drive)
+
+    def test_channel_build_inbound_messages_returns_when_already_user_closed(
+        self,
+    ):
+        channel = Channel(0, FakeConnection(), 360)
+        channel._user_closed = True
+        channel.set_state(channel.CLOSED)
+
+        messages = list(
+            channel.build_inbound_messages(break_on_empty=False),
+        )
+        self.assertEqual(messages, [])
+
+    def test_channel_build_inbound_messages_raises_when_closed_without_flag(
+        self,
+    ):
+        channel = Channel(0, FakeConnection(), 360)
+        channel.set_state(channel.CLOSED)
+
+        self.assertRaises(
+            AMQPChannelError,
+            lambda: list(
+                channel.build_inbound_messages(break_on_empty=False)
+            ),
+        )
+
+    def test_channel_user_initiated_close_helper(self):
+        connection = FakeConnection()
+        channel = Channel(0, connection, 360)
+
+        self.assertFalse(channel._user_initiated_close())
+
+        channel._user_closed = True
+        self.assertTrue(channel._user_initiated_close())
+
+        channel._user_closed = False
+        connection._user_closed = True
+        self.assertTrue(channel._user_initiated_close())
+
+    def test_channel_check_for_errors_still_raises_after_user_close(self):
+        connection = FakeConnection(state=FakeConnection.CLOSED)
+        connection._user_closed = True
+        channel = Channel(0, connection, 360)
+
+        self.assertRaises(AMQPConnectionError, channel.check_for_errors)
+
 
 class ChannelProcessDataEventTests(TestFramework):
     def test_channel_process_data_events(self):

amqpstorm/tests/unit/connection/test_connection.py

@@ -75,6 +75,14 @@ def test_connection_close_state(self):
 
         self.assertTrue(connection.is_closed)
 
+    def test_connection_close_sets_user_closed_flag(self):
+        connection = Connection('127.0.0.1', 'guest', 'guest', lazy=True)
+        connection.set_state(Connection.OPEN)
+
+        self.assertFalse(connection._user_closed)
+        connection.close()
+        self.assertTrue(connection._user_closed)
+
     def test_connection_open_channel_on_closed_connection(self):
         connection = Connection('127.0.0.1', 'guest', 'guest', lazy=True)