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.

1 files changed, 47 insertions, 4 deletions

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()