From b00f5d8ca7aa7a466f84fe92b605fbde350b7403 Mon Sep 17 00:00:00 2001
From: andy <github@anarkystic.com>
Date: Wed, 8 Sep 2010 22:43:54 +0200
Subject: [PATCH] Tests turn things into inlineCallbacks.

This is to duplicate the old behavior of BaseTestCase and to generally
prevent the bad situation in that tests appear to be passing when in
fact they haven't run because @defer.inlineCallbacks was forgotten.
---
 nova/rpc.py                  |  1 -
 nova/test.py                 | 51 +++++++++++++++++++++++++++++++++---
 nova/tests/cloud_unittest.py |  2 +-
 nova/tests/rpc_unittest.py   |  3 +--
 4 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/nova/rpc.py b/nova/rpc.py
index 7a89ca8d0..210d3cccf 100644
--- a/nova/rpc.py
+++ b/nova/rpc.py
@@ -250,7 +250,6 @@ def call(topic, msg):
     msg_id = uuid.uuid4().hex
     msg.update({'_msg_id': msg_id})
     LOG.debug("MSG_ID is %s" % (msg_id))
-
     conn = Connection.instance()
     d = defer.Deferred()
     consumer = DirectConsumer(connection=conn, msg_id=msg_id)
diff --git a/nova/test.py b/nova/test.py
index 5380d5743..1f4b33272 100644
--- a/nova/test.py
+++ b/nova/test.py
@@ -33,6 +33,7 @@ from twisted.trial import unittest
 
 from nova import fakerabbit
 from nova import flags
+from nova import rpc
 
 
 FLAGS = flags.FLAGS
@@ -63,22 +64,28 @@ class TrialTestCase(unittest.TestCase):
         self.stubs = stubout.StubOutForTesting()
         self.flag_overrides = {}
         self.injected = []
+        self._monkeyPatchAttach()
 
     def tearDown(self): # pylint: disable-msg=C0103
         """Runs after each test method to finalize/tear down test environment"""
-        super(TrialTestCase, self).tearDown()
         self.reset_flags()
         self.mox.UnsetStubs()
         self.stubs.UnsetAll()
         self.stubs.SmartUnsetAll()
         self.mox.VerifyAll()
-
+        
+        rpc.Consumer.attach_to_twisted = self.originalAttach
         for x in self.injected:
-            x.stop()
+            try:
+                x.stop()
+            except AssertionError:
+                pass
 
         if FLAGS.fake_rabbit:
             fakerabbit.reset_all()
 
+        super(TrialTestCase, self).tearDown()
+
     def flags(self, **kw):
         """Override flag variables for a test"""
         for k, v in kw.iteritems():
@@ -94,10 +101,46 @@ class TrialTestCase(unittest.TestCase):
         for k, v in self.flag_overrides.iteritems():
             setattr(FLAGS, k, v)
 
+    def run(self, result=None):
+        test_method = getattr(self, self._testMethodName)
+        setattr(self,
+                self._testMethodName,
+                self._maybeInlineCallbacks(test_method, result))
+        rv = super(TrialTestCase, self).run(result)
+        setattr(self, self._testMethodName, test_method)
+        return rv
+
+    def _maybeInlineCallbacks(self, func, result):
+        def _wrapped():
+            g = func()
+            if isinstance(g, defer.Deferred):
+                return g
+            if not hasattr(g, 'send'):
+                return defer.succeed(g)
+
+            inlined = defer.inlineCallbacks(func)
+            d = inlined()
+            return d
+        _wrapped.func_name = func.func_name
+        return _wrapped
+
+    def _monkeyPatchAttach(self):
+        self.originalAttach = rpc.Consumer.attach_to_twisted
+        def _wrapped(innerSelf):
+            rv = self.originalAttach(innerSelf)
+            self.injected.append(rv)
+            return rv
+
+        _wrapped.func_name = self.originalAttach.func_name
+        rpc.Consumer.attach_to_twisted = _wrapped
+
 
 class BaseTestCase(TrialTestCase):
     # TODO(jaypipes): Can this be moved into the TrialTestCase class?
-    """Base test case class for all unit tests."""
+    """Base test case class for all unit tests.
+    
+    DEPRECATED: This is being removed once Tornado is gone, use TrialTestCase.
+    """
     def setUp(self): # pylint: disable-msg=C0103
         """Run before each test method to initialize test environment"""
         super(BaseTestCase, self).setUp()
diff --git a/nova/tests/cloud_unittest.py b/nova/tests/cloud_unittest.py
index 0add53937..893717fe4 100644
--- a/nova/tests/cloud_unittest.py
+++ b/nova/tests/cloud_unittest.py
@@ -52,7 +52,7 @@ class CloudTestCase(test.TrialTestCase):
         self.compute_consumer = rpc.AdapterConsumer(connection=self.conn,
                                                     topic=FLAGS.compute_topic,
                                                     proxy=self.compute)
-        self.injected.append(self.compute_consumer.attach_to_twisted())
+        self.compute_consumer.attach_to_twisted()
 
         try:
             manager.AuthManager().create_user('admin', 'admin', 'admin')
diff --git a/nova/tests/rpc_unittest.py b/nova/tests/rpc_unittest.py
index e11967987..1bb2405c7 100644
--- a/nova/tests/rpc_unittest.py
+++ b/nova/tests/rpc_unittest.py
@@ -39,8 +39,7 @@ class RpcTestCase(test.TrialTestCase):
         self.consumer = rpc.AdapterConsumer(connection=self.conn,
                                             topic='test',
                                             proxy=self.receiver)
-
-        self.injected.append(self.consumer.attach_to_twisted())
+        self.consumer.attach_to_twisted()
 
     def test_call_succeed(self):
         """Get a value through rpc call"""