Add TestCase.stub_module to make stubbing modules easier

Previously various test cases would modify sys.modules directly,
duplicating code unnecessarily.

Change-Id: I09d333d6daf527e34db0262b28edcb0323cb2855

nova/test.py

@@ -24,6 +24,7 @@ inline callbacks.
 """
 
 import functools
+import sys
 import unittest
 import uuid
 
@@ -138,6 +139,7 @@ class TestCase(unittest.TestCase):
         self.stubs = stubout.StubOutForTesting()
         self.injected = []
         self._services = []
+        self._modules = {}
 
     def tearDown(self):
         """Runs after each test method to tear down test environment."""
@@ -151,6 +153,14 @@ class TestCase(unittest.TestCase):
             # Reset any overridden flags
             FLAGS.reset()
 
+            # Unstub modules
+            for name, mod in self._modules.iteritems():
+                if mod is not None:
+                    sys.modules[name] = mod
+                else:
+                    sys.modules.pop(name)
+            self._modules = {}
+
             # Stop any timers
             for x in self.injected:
                 try:
@@ -171,6 +181,11 @@ class TestCase(unittest.TestCase):
             for key in [k for k in self.__dict__.keys() if k[0] != '_']:
                 del self.__dict__[key]
 
+    def stub_module(self, name, mod):
+        if name not in self._modules:
+            self._modules[name] = sys.modules.get(name)
+        sys.modules[name] = mod
+
     def flags(self, **kw):
         """Override flag variables for a test."""
         for k, v in kw.iteritems():

nova/tests/xenapi/stubs.py

@@ -18,7 +18,6 @@
 import contextlib
 import pickle
 import random
-import sys
 
 from nova.openstack.common import jsonutils
 from nova import test
@@ -355,16 +354,6 @@ class XenAPITestBase(test.TestCase):
     def setUp(self):
         super(XenAPITestBase, self).setUp()
 
-        self.orig_XenAPI = sys.modules.get('XenAPI')
-        sys.modules['XenAPI'] = fake
+        self.stub_module('XenAPI', fake)
 
         fake.reset()
-
-    def tearDown(self):
-        if self.orig_XenAPI is not None:
-            sys.modules['XenAPI'] = self.orig_XenAPI
-            self.orig_XenAPI = None
-        else:
-            sys.modules.pop('XenAPI')
-
-        super(XenAPITestBase, self).tearDown()