X-Git-Url: https://gerrit.opnfv.org/gerrit/gitweb?a=blobdiff_plain;f=tests%2Funit%2Fcommon%2Ftest_openstack_utils.py;h=b3dc2d9c4a1551af5ed300b98bce4bbab6e4c9c1;hb=f73a21485dc963073f4232f1b131e503f457f9c5;hp=d610e181cb2fe4c28d90f069b6038bfb03a82f4c;hpb=f036e9898a69f5041f9cde02e3652c29e2de1643;p=yardstick.git

diff --git a/tests/unit/common/test_openstack_utils.py b/tests/unit/common/test_openstack_utils.py
index d610e181c..b3dc2d9c4 100644
--- a/tests/unit/common/test_openstack_utils.py
+++ b/tests/unit/common/test_openstack_utils.py
@@ -22,21 +22,25 @@ class GetCredentialsTestCase(unittest.TestCase):
 
     @mock.patch('yardstick.common.openstack_utils.os')
     def test_get_credentials(self, mock_os):
-        mock_os.getenv.return_value = ('2')
-        openstack_utils.get_credentials()
+        with mock.patch.dict('os.environ', {'OS_IDENTITY_API_VERSION': '2'},
+                             clear=True):
+            openstack_utils.get_credentials()
 
 
 class GetHeatApiVersionTestCase(unittest.TestCase):
 
-    @mock.patch('yardstick.common.openstack_utils.os')
-    def test_get_heat_api_version(self, mock_os):
+    def test_get_heat_api_version_check_result(self):
         API = 'HEAT_API_VERSION'
-        openstack_utils.get_heat_api_version()
-        mock_os.getenv.assert_called_with(API)
-
-    @mock.patch('yardstick.common.openstack_utils.os')
-    def test_get_heat_api_version(self, mock_os):
-        mock_os.getenv.return_value = ('2')
         expected_result = '2'
-        api_version = openstack_utils.get_heat_api_version()
-        self.assertEqual(api_version, expected_result)
+
+        with mock.patch.dict('os.environ', {API: '2'}, clear=True):
+            api_version = openstack_utils.get_heat_api_version()
+            self.assertEqual(api_version, expected_result)
+
+
+def main():
+    unittest.main()
+
+
+if __name__ == '__main__':
+    main()