Support using existing private key when using external heat template file

[yardstick.git] / yardstick / tests / unit / service / test_environment.py

diff --git a/yardstick/tests/unit/service/test_environment.py b/yardstick/tests/unit/service/test_environment.py
index 4af9a39..779e6ea 100644 (file)
--- a/yardstick/tests/unit/service/test_environment.py
+++ b/yardstick/tests/unit/service/test_environment.py
@@ -6,16 +6,15 @@
 # which accompanies this distribution, and is available at
 # http://www.apache.org/licenses/LICENSE-2.0
 ##############################################################################
-import unittest
 
 import mock
 
-from yardstick.service.environment import Environment
-from yardstick.service.environment import AnsibleCommon
-from yardstick.common.exceptions import UnsupportedPodFormatError
+from yardstick.common import exceptions
+from yardstick.service import environment
+from yardstick.tests.unit import base as ut_base
 
 
-class EnvironmentTestCase(unittest.TestCase):
+class EnvironmentTestCase(ut_base.BaseUnitTestCase):
 
     def test_get_sut_info(self):
         pod_info = {
@@ -31,19 +30,17 @@ class EnvironmentTestCase(unittest.TestCase):
             ]
         }
 
-        AnsibleCommon.gen_inventory_ini_dict = mock.MagicMock()
-        AnsibleCommon.get_sut_info = mock.MagicMock(return_value={'node1': {}})
-
-        env = Environment(pod=pod_info)
-        env.get_sut_info()
+        with mock.patch.object(environment.AnsibleCommon,
+                               'gen_inventory_ini_dict'), \
+                mock.patch.object(environment.AnsibleCommon, 'get_sut_info',
+                                  return_value={'node1': {}}), \
+                mock.patch.object(environment.Environment, '_format_sut_info'):
+            env = environment.Environment(pod=pod_info)
+            env.get_sut_info()
 
     def test_get_sut_info_pod_str(self):
         pod_info = 'nodes'
 
-        env = Environment(pod=pod_info)
-        with self.assertRaises(UnsupportedPodFormatError):
+        env = environment.Environment(pod=pod_info)
+        with self.assertRaises(exceptions.UnsupportedPodFormatError):
             env.get_sut_info()
-
-
-if __name__ == '__main__':
-    unittest.main()