Repository: pika/pika

Issue:
Pass connection instance as first arg of callbacks for add_on_connection_blocked_callback and add_on_connection_unblocked_callback
Targeting 1.0.0

Discussion:
vitaly-krugl: Ready for review

Patch:
diff --git a/pika/adapters/blocking_connection.py b/pika/adapters/blocking_connection.py
index b1fb7dc..be1db63 100644
--- a/pika/adapters/blocking_connection.py
+++ b/pika/adapters/blocking_connection.py
@@ -542,22 +542,24 @@ class BlockingConnection(object):
         self.add_timeout(0, user_callback)
 
 
-    def _on_connection_blocked(self, user_callback, method_frame):
+    def _on_connection_blocked(self, user_callback, _impl, method_frame):
         """Handle Connection.Blocked notification from RabbitMQ broker
 
         :param callable user_callback: callback passed to
            `add_on_connection_blocked_callback`
+        :param SelectConnection _impl:
         :param pika.frame.Method method_frame: method frame having `method`
             member of type `pika.spec.Connection.Blocked`
         """
         self._ready_events.append(
             _ConnectionBlockedEvt(user_callback, method_frame))
 
-    def _on_connection_unblocked(self, user_callback, method_frame):
+    def _on_connection_unblocked(self, user_callback, _impl, method_frame):
         """Handle Connection.Unblocked notification from RabbitMQ broker
 
         :param callable user_callback: callback passed to
            `add_on_connection_unblocked_callback`
+        :param SelectConnection _impl:
         :param pika.frame.Method method_frame: method frame having `method`
             member of type `pika.spec.Connection.Blocked`
         """
@@ -595,13 +597,14 @@ class BlockingConnection(object):
         See also `ConnectionParameters.blocked_connection_timeout`.
 
         :param method callback: Callback to call on `Connection.Blocked`,
-            having the signature `callback(pika.frame.Method)`, where the
-            method frame's `method` member is of type
-            `pika.spec.Connection.Blocked`
+            having the signature `callback(connection, pika.frame.Method)`,
+            where connection is the `BlockingConnection` instance and the method
+            frame's `method` member is of type `pika.spec.Connection.Blocked`
 
         """
         self._impl.add_on_connection_blocked_callback(
-            functools.partial(self._on_connection_blocked, callback))
+            functools.partial(self._on_connection_blocked,
+                              functools.partial(callback, self)))
 
     def add_on_connection_unblocked_callback(self, callback):
         """Add a callback to be notified when RabbitMQ has sent a
@@ -609,14 +612,15 @@ class BlockingConnection(object):
         to start publishing again. The callback will be passed the
         `Connection.Unblocked` method frame.
 
-        :param method callback: Callback to call on
-            `Connection.Unblocked`, having the signature
-            `callback(pika.frame.Method)`, where the method frame's
-            `method` member is of type `pika.spec.Connection.Unblocked`
+        :param method callback: Callback to call on Connection.Unblocked`,
+            having the signature `callback(connection, pika.frame.Method)`,
+            where connection is the `BlockingConnection` instance and the method
+             frame's `method` member is of type `pika.spec.Connection.Unblocked`
 
         """
         self._impl.add_on_connection_unblocked_callback(
-            functools.partial(self._on_connection_unblocked, callback))
+            functools.partial(self._on_connection_unblocked,
+                              functools.partial(callback, self)))
 
     def add_timeout(self, deadline, callback):
         """Create a single-shot timer to fire after deadline seconds. Do not
diff --git a/pika/connection.py b/pika/connection.py
index c784cc2..25ec506 100644
--- a/pika/connection.py
+++ b/pika/connection.py
@@ -6,6 +6,7 @@ import ast
 import sys
 import collections
 import copy
+import functools
 import logging
 import math
 import numbers
@@ -1018,11 +1019,9 @@ class Connection(object):
     # pylint: disable=R0201
 
     ON_CONNECTION_BACKPRESSURE = '_on_connection_backpressure'
-    ON_CONNECTION_BLOCKED = '_on_connection_blocked'
     ON_CONNECTION_CLOSED = '_on_connection_closed'
     ON_CONNECTION_ERROR = '_on_connection_error'
     ON_CONNECTION_OPEN = '_on_connection_open'
-    ON_CONNECTION_UNBLOCKED = '_on_connection_unblocked'
     CONNECTION_CLOSED = 0
     CONNECTION_INIT = 1
     CONNECTION_PROTOCOL = 2
@@ -1146,19 +1145,21 @@ class Connection(object):
         instead of relying on back pressure throttling. The callback
         will be passed the ``Connection.Blocked`` method frame.
 
-        TODO Also pass the connection as the callback's first arg
-
         See also `ConnectionParameters.blocked_connection_timeout`.
 
         :param method callback: Callback to call on `Connection.Blocked`,
-            having the signature `callback(pika.frame.Method)`, where the
-            method frame's `method` member is of type
+            having the signature `callback(connection, pika.frame.Method)`,
+            where the method frame's `method` member is of type
             `pika.spec.Connection.Blocked`
 
         """
         if not callable(callback):
             raise TypeError('callback should be a function or method.')
-        self.callbacks.add(0, spec.Connection.Blocked, callback, False)
+
+        self.callbacks.add(0,
+                           spec.Connection.Blocked,
+                           functools.partial(callback, self),
+                           one_shot=False)
 
     def add_on_connection_unblocked_callback(self, callback):
         """Add a callback to be notified when RabbitMQ has sent a
@@ -1166,17 +1167,19 @@ class Connection(object):
         to start publishing again. The callback will be passed the
         ``Connection.Unblocked`` method frame.
 
-        TODO Also pass the connection as the callback's first arg
-
         :param method callback: Callback to call on
             `Connection.Unblocked`, having the signature
-            `callback(pika.frame.Method)`, where the method frame's
+            `callback(connection, pika.frame.Method)`, where the method frame's
             `method` member is of type `pika.spec.Connection.Unblocked`
 
         """
         if not callable(callback):
             raise TypeError('callback should be a function or method.')
-        self.callbacks.add(0, spec.Connection.Unblocked, callback, False)
+
+        self.callbacks.add(0,
+                           spec.Connection.Unblocked,
+                           functools.partial(callback, self),
+                           one_shot=False)
 
     def add_on_open_callback(self, callback):
         """Add a callback notification when the connection has opened. The
@@ -1789,7 +1792,7 @@ class Connection(object):
         self._on_terminate(InternalCloseReasons.BLOCKED_CONNECTION_TIMEOUT,
                            'Blocked connection timeout expired')
 
-    def _on_connection_blocked(self, method_frame):
+    def _on_connection_blocked(self, _connection, method_frame):
         """Handle Connection.Blocked notification from RabbitMQ broker
 
         :param pika.frame.Method method_frame: method frame having `method`
@@ -1808,7 +1811,7 @@ class Connection(object):
                 self.params.blocked_connection_timeout,
                 self._on_blocked_connection_timeout)
 
-    def _on_connection_unblocked(self, method_frame):
+    def _on_connection_unblocked(self, _connection, method_frame):
         """Handle Connection.Unblocked notification from RabbitMQ broker
 
         :param pika.frame.Method method_frame: method frame having `method`


Test patch:
diff --git a/tests/acceptance/async_adapter_tests.py b/tests/acceptance/async_adapter_tests.py
index d27a625..0539596 100644
--- a/tests/acceptance/async_adapter_tests.py
+++ b/tests/acceptance/async_adapter_tests.py
@@ -445,9 +445,11 @@ class TestBlockedConnectionTimesOut(AsyncTestCase, AsyncAdapters):  # pylint: di
     def begin(self, channel):
 
         # Simulate Connection.Blocked
-        channel.connection._on_connection_blocked(pika.frame.Method(
-            0,
-            pika.spec.Connection.Blocked('Testing blocked connection timeout')))
+        channel.connection._on_connection_blocked(
+            channel.connection,
+            pika.frame.Method(0,
+                              pika.spec.Connection.Blocked(
+                                  'Testing blocked connection timeout')))
 
     def on_closed(self, connection, reply_code, reply_text):
         """called when the connection has finished closing"""
@@ -471,15 +473,16 @@ class TestBlockedConnectionUnblocks(AsyncTestCase, AsyncAdapters):  # pylint: di
     def begin(self, channel):
 
         # Simulate Connection.Blocked
-        channel.connection._on_connection_blocked(pika.frame.Method(
-            0,
-            pika.spec.Connection.Blocked(
-                'Testing blocked connection unblocks')))
+        channel.connection._on_connection_blocked(
+            channel.connection,
+            pika.frame.Method(0,
+                              pika.spec.Connection.Blocked(
+                                  'Testing blocked connection unblocks')))
 
         # Simulate Connection.Unblocked
-        channel.connection._on_connection_unblocked(pika.frame.Method(
-            0,
-            pika.spec.Connection.Unblocked()))
+        channel.connection._on_connection_unblocked(
+            channel.connection,
+            pika.frame.Method(0, pika.spec.Connection.Unblocked()))
 
         # Schedule shutdown after blocked connection timeout would expire
         channel.connection.add_timeout(0.005, self.on_cleanup_timer)
diff --git a/tests/acceptance/blocking_adapter_test.py b/tests/acceptance/blocking_adapter_test.py
index 5623922..03b4681 100644
--- a/tests/acceptance/blocking_adapter_test.py
+++ b/tests/acceptance/blocking_adapter_test.py
@@ -401,27 +401,34 @@ class TestConnectionRegisterForBlockAndUnblock(BlockingTestCaseBase):
 
         # NOTE: I haven't figured out yet how to coerce RabbitMQ to emit
         # Connection.Block and Connection.Unblock from the test, so we'll
-        # just call the registration functions for now, to make sure that
-        # registration doesn't crash
-
-        connection.add_on_connection_blocked_callback(lambda frame: None)
+        # just call the registration functions for now and simulate incoming
+        # blocked/unblocked frames
 
         blocked_buffer = []
-        evt = blocking_connection._ConnectionBlockedEvt(
-            lambda f: blocked_buffer.append("blocked"),
-            pika.frame.Method(1, pika.spec.Connection.Blocked('reason')))
-        repr(evt)
-        evt.dispatch()
-        self.assertEqual(blocked_buffer, ["blocked"])
+        connection.add_on_connection_blocked_callback(
+            lambda conn, frame: blocked_buffer.append((conn, frame)))
+        # Simulate dispatch of blocked connection
+        blocked_frame = pika.frame.Method(
+            0,
+            pika.spec.Connection.Blocked('reason'))
+        connection._impl._process_frame(blocked_frame)
+        connection.sleep(0) # facilitate dispatch of pending events
+        self.assertEqual(len(blocked_buffer), 1)
+        conn, frame = blocked_buffer[0]
+        self.assertIs(conn, connection)
+        self.assertIs(frame, blocked_frame)
 
         unblocked_buffer = []
-        connection.add_on_connection_unblocked_callback(lambda frame: None)
-        evt = blocking_connection._ConnectionUnblockedEvt(
-            lambda f: unblocked_buffer.append("unblocked"),
-            pika.frame.Method(1, pika.spec.Connection.Unblocked()))
-        repr(evt)
-        evt.dispatch()
-        self.assertEqual(unblocked_buffer, ["unblocked"])
+        connection.add_on_connection_unblocked_callback(
+            lambda conn, frame: unblocked_buffer.append((conn, frame)))
+        # Simulate dispatch of unblocked connection
+        unblocked_frame = pika.frame.Method(0, pika.spec.Connection.Unblocked())
+        connection._impl._process_frame(unblocked_frame)
+        connection.sleep(0) # facilitate dispatch of pending events
+        self.assertEqual(len(unblocked_buffer), 1)
+        conn, frame = unblocked_buffer[0]
+        self.assertIs(conn, connection)
+        self.assertIs(frame, unblocked_frame)
 
 
 class TestBlockedConnectionTimeout(BlockingTestCaseBase):
@@ -436,9 +443,11 @@ class TestBlockedConnectionTimeout(BlockingTestCaseBase):
         # simulate it for now
 
         # Simulate Connection.Blocked
-        conn._impl._on_connection_blocked(pika.frame.Method(
-            0,
-            pika.spec.Connection.Blocked('TestBlockedConnectionTimeout')))
+        conn._impl._on_connection_blocked(
+            conn._impl,
+            pika.frame.Method(
+                0,
+                pika.spec.Connection.Blocked('TestBlockedConnectionTimeout')))
 
         # Wait for connection teardown
         with self.assertRaises(pika.exceptions.ConnectionClosed) as excCtx:
diff --git a/tests/unit/blocking_connection_tests.py b/tests/unit/blocking_connection_tests.py
index f303023..5fda321 100644
--- a/tests/unit/blocking_connection_tests.py
+++ b/tests/unit/blocking_connection_tests.py
@@ -15,6 +15,19 @@ import pika.channel
 from pika.exceptions import AMQPConnectionError, ChannelClosed
 
 
+# Disable protected-access
+# pylint: disable=W0212
+
+# Disable missing-docstring
+# pylint: disable=C0111
+
+# Disable invalid-name
+# pylint: disable=C0103
+
+# Disable no-self-use
+# pylint: disable=R0201
+
+
 class BlockingConnectionMockTemplate(blocking_connection.BlockingConnection):
     pass
 
@@ -325,3 +338,25 @@ class BlockingConnectionTests(unittest.TestCase):
             # and the raised error must then looks like:
             self.assertEqual('Connection to 127.0.0.1:5672 failed: timeout',
                              str(ctx.exception))
+
+    def test_connection_blocked_evt(self):
+        blocked_buffer = []
+        frame = pika.frame.Method(0, pika.spec.Connection.Blocked('reason'))
+        evt = blocking_connection._ConnectionBlockedEvt(
+            blocked_buffer.append,
+            frame)
+        repr(evt)
+        evt.dispatch()
+        self.assertEqual(len(blocked_buffer), 1)
+        self.assertIs(blocked_buffer[0], frame)
+
+    def test_connection_unblocked_evt(self):
+        unblocked_buffer = []
+        frame = pika.frame.Method(0, pika.spec.Connection.Unblocked())
+        evt = blocking_connection._ConnectionUnblockedEvt(
+            unblocked_buffer.append,
+            frame)
+        repr(evt)
+        evt.dispatch()
+        self.assertEqual(len(unblocked_buffer), 1)
+        self.assertIs(unblocked_buffer[0], frame)
diff --git a/tests/unit/connection_tests.py b/tests/unit/connection_tests.py
index 38db1d7..9d99314 100644
--- a/tests/unit/connection_tests.py
+++ b/tests/unit/connection_tests.py
@@ -678,6 +678,36 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
             if frame_type == frame.Heartbeat:
                 self.assertTrue(self.connection.heartbeat.received.called)
 
+    def test_add_on_connection_blocked_callback(self):
+        blocked_buffer = []
+        self.connection.add_on_connection_blocked_callback(
+            lambda conn, frame: blocked_buffer.append((conn, frame)))
+
+        # Simulate dispatch of blocked connection
+        blocked_frame = pika.frame.Method(
+            0,
+            pika.spec.Connection.Blocked('reason'))
+        self.connection._process_frame(blocked_frame)
+
+        self.assertEqual(len(blocked_buffer), 1)
+        conn, frame = blocked_buffer[0]
+        self.assertIs(conn, self.connection)
+        self.assertIs(frame, blocked_frame)
+
+    def test_add_on_connection_unblocked_callback(self):
+        unblocked_buffer = []
+        self.connection.add_on_connection_unblocked_callback(
+            lambda conn, frame: unblocked_buffer.append((conn, frame)))
+
+        # Simulate dispatch of unblocked connection
+        unblocked_frame = pika.frame.Method(0, pika.spec.Connection.Unblocked())
+        self.connection._process_frame(unblocked_frame)
+
+        self.assertEqual(len(unblocked_buffer), 1)
+        conn, frame = unblocked_buffer[0]
+        self.assertIs(conn, self.connection)
+        self.assertIs(frame, unblocked_frame)
+
     @mock.patch.object(
         connection.Connection,
         'connect',
@@ -714,6 +744,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
                 blocked_connection_timeout=60))
 
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         # Check
@@ -736,6 +767,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
 
         # Simulate Connection.Blocked trigger
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         # Check
@@ -748,6 +780,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
 
         # Simulate Connection.Blocked trigger again
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         self.assertEqual(conn.add_timeout.call_count, 1)
@@ -770,6 +803,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
                 blocked_connection_timeout=60))
 
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         conn._on_blocked_connection_timeout()
@@ -798,6 +832,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
                 blocked_connection_timeout=60))
 
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         self.assertIsNotNone(conn._blocked_conn_timer)
@@ -805,6 +840,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
         timer = conn._blocked_conn_timer
 
         conn._on_connection_unblocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Unblocked)'))
 
         # Check
@@ -829,6 +865,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
 
         # Simulate Connection.Blocked
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         self.assertIsNotNone(conn._blocked_conn_timer)
@@ -837,6 +874,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
 
         # Simulate Connection.Unblocked
         conn._on_connection_unblocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Unblocked)'))
 
         # Check
@@ -845,6 +883,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
 
         # Simulate Connection.Unblocked again
         conn._on_connection_unblocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Unblocked)'))
 
         self.assertEqual(conn.remove_timeout.call_count, 1)
@@ -872,6 +911,7 @@ class ConnectionTests(unittest.TestCase):  # pylint: disable=R0904
                 blocked_connection_timeout=60))
 
         conn._on_connection_blocked(
+            conn,
             mock.Mock(name='frame.Method(Connection.Blocked)'))
 
         self.assertIsNotNone(conn._blocked_conn_timer)