Merge trunk.

1 files changed, 49 insertions, 5 deletions

diff --git a/nova/test.py b/nova/test.py
index 5ed0c73d3..08e1dea2d 100644
--- a/nova/test.py
+++ b/nova/test.py
@@ -34,6 +34,7 @@ from twisted.trial import unittest
 from nova import db
 from nova import fakerabbit
 from nova import flags
+from nova import rpc
 
 
 FLAGS = flags.FLAGS
@@ -63,20 +64,30 @@ class TrialTestCase(unittest.TestCase):
         self.mox = mox.Mox()
         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:
+            try:
+                x.stop()
+            except AssertionError:
+                pass
 
         if FLAGS.fake_rabbit:
             fakerabbit.reset_all()
         db.security_group_destroy_all(None)
 
+        super(TrialTestCase, self).tearDown()
+
     def flags(self, **kw):
         """Override flag variables for a test"""
         for k, v in kw.iteritems():
@@ -92,16 +103,51 @@ 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()
         # TODO(termie): we could possibly keep a more global registry of
         #               the injected listeners... this is fine for now though
-        self.injected = []
         self.ioloop = ioloop.IOLoop.instance()
 
         self._waiting = None
@@ -111,8 +157,6 @@ class BaseTestCase(TrialTestCase):
     def tearDown(self):# pylint: disable-msg=C0103
         """Runs after each test method to finalize/tear down test environment"""
         super(BaseTestCase, self).tearDown()
-        for x in self.injected:
-            x.stop()
         if FLAGS.fake_rabbit:
             fakerabbit.reset_all()