1 files changed, 9 insertions, 12 deletions
diff --git a/yardstick/tests/unit/benchmark/contexts/test_heat.py b/yardstick/tests/unit/benchmark/contexts/test_heat.py
index ebb1d69ba..9c822b3a7 100644
--- a/yardstick/tests/unit/benchmark/contexts/test_heat.py
+++ b/yardstick/tests/unit/benchmark/contexts/test_heat.py
@@ -8,38 +8,34 @@
 ##############################################################################
 
 from collections import OrderedDict
-from itertools import count
 import logging
 import os
 
 import mock
-import unittest
 
+from yardstick import ssh
 from yardstick.benchmark.contexts import base
 from yardstick.benchmark.contexts import heat
 from yardstick.benchmark.contexts import model
 from yardstick.common import constants as consts
 from yardstick.common import exceptions as y_exc
-from yardstick import ssh
+from yardstick.tests.unit import base as ut_base
 
 
 LOG = logging.getLogger(__name__)
 
 
-class HeatContextTestCase(unittest.TestCase):
-
-    def __init__(self, *args, **kwargs):
-        super(HeatContextTestCase, self).__init__(*args, **kwargs)
-        self.name_iter = ('vnf{:03}'.format(x) for x in count(0, step=3))
+class HeatContextTestCase(ut_base.BaseUnitTestCase):
 
     def setUp(self):
         self.test_context = heat.HeatContext()
         self.addCleanup(self._remove_contexts)
-        self.mock_context = mock.Mock(spec=heat.HeatContext())
 
-    def _remove_contexts(self):
-        if self.test_context in self.test_context.list:
-            self.test_context._delete_context()
+    @staticmethod
+    def _remove_contexts():
+        for context in base.Context.list:
+            context._delete_context()
+        base.Context.list = []
 
     def test___init__(self):
         self.assertIsNone(self.test_context._name)
@@ -658,6 +654,7 @@ class HeatContextTestCase(unittest.TestCase):
         baz3_server.public_ip = None
         baz3_server.context.user = 'zab'
 
+        self.mock_context = mock.Mock(spec=heat.HeatContext())
         self.mock_context._name = 'bar1'
         self.test_context.stack = mock.Mock()
         self.mock_context.stack.outputs = {