refactor: add reference to aiostomp instance

pedrokiefer · Feb 6, 2019 · 48d422e · 48d422e
1 parent 716e081
commit 48d422e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/aiostomp/aiostomp.py b/aiostomp/aiostomp.py
@@ -160,7 +160,7 @@ async def _reconnect(self):
                 else:
                     logger.error('All connections attempts failed.')
                     if self._on_error:
-                        asyncio.ensure_future(self._on_error(ExceededRetryCount()), loop=self._loop)
+                        asyncio.ensure_future(self._on_error(ExceededRetryCount(self)), loop=self._loop)
                     break
 
                 self._increment_retry_interval()

diff --git a/aiostomp/errors.py b/aiostomp/errors.py
@@ -11,4 +11,7 @@ class StompDisconnectedError(Exception):
 
 
 class ExceededRetryCount(Exception):
-    pass
+
+    def __init__(self, ref):
+        super().__init__('Retry count exceeded!')
+        self.ref = ref
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -534,6 +534,7 @@ async def test_reconnection_error(self, logger_mock):
         logger_mock.error.assert_called_with(
             'All connections attempts failed.')
         self.assertIsInstance(self.stomp._on_error.call_args[0][0], ExceededRetryCount)
+        self.assertEqual(self.stomp._on_error.call_args[0][0].ref, self.stomp)
 
     @unittest_run_loop
     async def test_can_reconnect_on_connection_lost(self):