X-Git-Url: https://gerrit.opnfv.org/gerrit/gitweb?a=blobdiff_plain;ds=sidebyside;f=yardstick%2Ftests%2Funit%2Fservice%2Ftest_environment.py;h=779e6eaa03d212e3113b3dc0336a8644ac7d7bf0;hb=0ff5b131fa3f93f8d47b46c2f13661c5b88d37a1;hp=4af9a3958caa0c7a73539e15b537a58f46c3927b;hpb=c3f91fac56cd938c408c62ff44f5ff039b1fd8e6;p=yardstick.git

diff --git a/yardstick/tests/unit/service/test_environment.py b/yardstick/tests/unit/service/test_environment.py
index 4af9a3958..779e6eaa0 100644
--- 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()