Add TestCase.stub_module to make stubbing modules easier

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

Change-Id: I09d333d6daf527e34db0262b28edcb0323cb2855

3 files changed, 17 insertions, 15 deletions

diff --git a/nova/test.py b/nova/test.py
index 515ba48ef..7fb01e9ab 100644
--- a/nova/test.py
+++ b/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():
diff --git a/nova/tests/network/test_manager.py b/nova/tests/network/test_manager.py
index 2f597547c..c24bee2d8 100644
--- a/nova/tests/network/test_manager.py
+++ b/nova/tests/network/test_manager.py
@@ -1846,8 +1846,7 @@ class LdapDNSTestCase(test.TestCase):
     def setUp(self):
         super(LdapDNSTestCase, self).setUp()
 
-        self.saved_ldap = sys.modules.get('ldap')
-        sys.modules['ldap'] = fake_ldap
+        self.stub_module('ldap', fake_ldap)
 
         self.driver = FakeLdapDNS()
         self.driver.create_domain(domain1)
@@ -1856,7 +1855,6 @@ class LdapDNSTestCase(test.TestCase):
     def tearDown(self):
         self.driver.delete_domain(domain1)
         self.driver.delete_domain(domain2)
-        sys.modules['ldap'] = self.saved_ldap
         super(LdapDNSTestCase, self).tearDown()
 
     def test_ldap_dns_domains(self):
diff --git a/nova/tests/xenapi/stubs.py b/nova/tests/xenapi/stubs.py
index 35aa8e4aa..cde17a7af 100644
--- a/nova/tests/xenapi/stubs.py
+++ b/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()