Merge pull request #7 from OpenMatchmaking/feature-custom-event-loop

Specifying a custom event loop in `run` method

Author: Relrin

sage_utils/__init__.py

@@ -1,4 +1,4 @@
 __title__ = 'sage-utils'
-__version__ = '0.5.1'
+__version__ = '0.5.2'
 __license__ = 'BSD'
 VERSION = __version__

sage_utils/amqp/workers.py

@@ -26,7 +26,8 @@ def get_microservice_data(self, app):
         raise NotImplementedError("The `get_microservice_data(data)` method "
                                   "must be implemented.")
 
-    async def run(self, *args, **kwargs):
+    async def run(self, *args, loop=None, **kwargs):
+        loop = loop or getattr(self.app, 'loop', None) or asyncio.get_event_loop()
         client = RpcAmqpClient(
             self.app,
             routing_key=self.REQUEST_QUEUE_NAME,
@@ -58,7 +59,9 @@ async def run(self, *args, **kwargs):
                     "responding or the required queues and the exchanges aren't created. "
                     "Retry after {} second(s).".format(self.RETRY_TIMEOUT)
                 )
-                await asyncio.sleep(self.RETRY_TIMEOUT, loop=self.app.loop)
+                await asyncio.sleep(self.RETRY_TIMEOUT, loop=loop)
+
+        await self.free_resources()
 
         if not is_registered:
             raise ConnectionError('Occurred an error during registering microservice.')

tests/amqp/test_worker.py

@@ -51,6 +51,30 @@ async def test_microservice_has_been_registered_successfully(event_loop):
     await extension.deinit(event_loop)
 
 
+@pytest.mark.asyncio
+async def test_microservice_has_been_registered_successfully_with_custom_event_loop(event_loop):
+    app = Application(config=FakeConfig(), loop=event_loop)
+    register_worker = FakeRegisterMicroserviceWorker(app)
+    extension = AmqpExtension(app)
+    extension.register_worker(register_worker)
+
+    await extension.init(event_loop)
+
+    instance = BaseRegisterWorker(app)
+    instance.REQUEST_QUEUE_NAME = REQUEST_QUEUE
+    instance.REQUEST_EXCHANGE_NAME = REQUEST_EXCHANGE
+    instance.RESPONSE_EXCHANGE_NAME = RESPONSE_EXCHANGE_NAME
+
+    with mock.patch(
+        target='sage_utils.amqp.workers.BaseRegisterWorker.get_microservice_data',
+        side_effect=valid_microservice_data
+    ):
+        result = await instance.run(loop=event_loop)
+        assert result is None
+
+    await extension.deinit(event_loop)
+
+
 @pytest.mark.asyncio
 async def test_register_worker_raises_an_error_with_invalid_response(event_loop):
     app = Application(config=FakeConfig(), loop=event_loop)