From b0007676fb817eb0b98554d27789cdaccd2c72a7 Mon Sep 17 00:00:00 2001
From: nico <nico@beerfactory.org>
Date: Fri, 21 Aug 2015 13:08:52 +0200
Subject: [PATCH] map() method now returns a map of result {plugin1: ret1,
 plugin2: ret2...} for each plugin called

---
 hbmqtt/plugins/manager.py     |  9 +++++++--
 tests/plugins/test_manager.py | 14 ++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/hbmqtt/plugins/manager.py b/hbmqtt/plugins/manager.py
index 89a044f..af42bda 100644
--- a/hbmqtt/plugins/manager.py
+++ b/hbmqtt/plugins/manager.py
@@ -146,12 +146,17 @@ class PluginManager:
         :return:
         """
         tasks = []
+        plugins_list = []
         for plugin in self._plugins:
             coro_instance = coro(plugin, *args, **kwargs)
             if coro_instance:
                 tasks.append(self._schedule_coro(coro_instance))
-        ret = yield from asyncio.gather(*tasks, loop=self._loop)
-        return ret
+                plugins_list.append(plugin)
+        ret_list = yield from asyncio.gather(*tasks, loop=self._loop)
+
+        # Create result map plugin=>ret
+        ret_dict = {k: v for k, v in zip(plugins_list, ret_list)}
+        return ret_dict
 
     @staticmethod
     def _get_coro(plugin, coro_name, *args, **kwargs):
diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py
index 2194c0d..d18394e 100644
--- a/tests/plugins/test_manager.py
+++ b/tests/plugins/test_manager.py
@@ -30,6 +30,10 @@ class EventTestPlugin:
     def test_coro(self):
         self.coro_flag = True
 
+    @asyncio.coroutine
+    def ret_coro(self):
+        return "TEST"
+
 
 class TestPluginManager(unittest.TestCase):
     def setUp(self):
@@ -71,3 +75,13 @@ class TestPluginManager(unittest.TestCase):
         self.loop.run_until_complete(call_coro())
         plugin = manager.get_plugin("event_plugin")
         self.assertTrue(plugin.object.test_coro)
+
+    def test_map_coro_return(self):
+        @asyncio.coroutine
+        def call_coro():
+            return (yield from manager.map_plugin_coro('ret_coro'))
+
+        manager = PluginManager("hbmqtt.test.plugins", context=None, loop=self.loop)
+        ret = self.loop.run_until_complete(call_coro())
+        plugin = manager.get_plugin("event_plugin")
+        self.assertEqual(ret[plugin], "TEST")