From fa547237aea3e27fc4eaab4fb2cc144ad2d0466e Mon Sep 17 00:00:00 2001
From: termie <github@anarkystic.com>
Date: Wed, 25 May 2011 15:42:25 -0700
Subject: [PATCH] make sure that using multicall on a call with a single result
 still functions

---
 nova/rpc.py            |  4 ++--
 nova/tests/test_rpc.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/nova/rpc.py b/nova/rpc.py
index 7faed4d3..84493271 100644
--- a/nova/rpc.py
+++ b/nova/rpc.py
@@ -233,10 +233,10 @@ class AdapterConsumer(Consumer):
                     logging.error('rval! %s', rval)
                     for x in rval:
                         msg_reply(msg_id, x, None)
-                    msg_reply(msg_id, None, None)
                 else:
                     msg_reply(msg_id, rval, None)
-                #msg_reply(msg_id, rval, None)
+                # This final None tells multicall that it is done.
+                msg_reply(msg_id, None, None)
         except Exception as e:
             logging.exception('Exception during message handling')
             if msg_id:
diff --git a/nova/tests/test_rpc.py b/nova/tests/test_rpc.py
index e5d99474..c1ef60ff 100644
--- a/nova/tests/test_rpc.py
+++ b/nova/tests/test_rpc.py
@@ -49,6 +49,35 @@ class RpcTestCase(test.TestCase):
                                                  "args": {"value": value}})
         self.assertEqual(value, result)
 
+    def test_call_succeed_despite_multiple_returns(self):
+        """Get a value through rpc call"""
+        value = 42
+        result = rpc.call(self.context, 'test', {"method": "echo_three_times",
+                                                 "args": {"value": value}})
+        self.assertEqual(value, result)
+
+    def test_call_succeed_despite_multiple_returns_yield(self):
+        """Get a value through rpc call"""
+        value = 42
+        result = rpc.call(self.context, 'test',
+                          {"method": "echo_three_times_yield",
+                           "args": {"value": value}})
+        self.assertEqual(value, result)
+
+    def test_multicall_succeed_once(self):
+        """Get a value through rpc call"""
+        value = 42
+        result = rpc.multicall(self.context,
+                              'test',
+                              {"method": "echo",
+                               "args": {"value": value}})
+        i = 0
+        for x in result:
+            if i > 0:
+                self.fail('should only receive one response')
+            self.assertEqual(value + i, x)
+            i += 1
+
     def test_multicall_succeed_three_times(self):
         """Get a value through rpc call"""
         value = 42