diff --git a/qemu/tests/balloon_check.py b/qemu/tests/balloon_check.py
index 10e3a4edcd469e5b66291c50680ef102b1a38850..70fb28467c077ab0d63f13ed2ca2bd8d868d4712 100644
--- a/qemu/tests/balloon_check.py
+++ b/qemu/tests/balloon_check.py
@@ -2,10 +2,13 @@ import time
 import re
 import logging
 import random
-from autotest.client.shared import error
+
+from avocado.core import exceptions
+
 from virttest import qemu_monitor
 from virttest import utils_test
 from virttest import utils_misc
+from virttest import error_context
 from virttest.utils_test.qemu import MemoryBaseTest
 
 
@@ -49,7 +52,7 @@ class BallooningTest(MemoryBaseTest):
             return 0
         return ballooned_mem
 
-    @error.context_aware
+    @error_context.context_aware
     def memory_check(self, step, ballooned_mem):
         """
         Check memory status according expect values
@@ -61,7 +64,7 @@ class BallooningTest(MemoryBaseTest):
         :return: memory size get from monitor and guest
         :rtype: tuple
         """
-        error.context("Check memory status %s" % step, logging.info)
+        error_context.context("Check memory status %s" % step, logging.info)
         mmem = self.get_ballooned_memory()
         gmem = self.get_memory_status()
         # for windows illegal test:set windows guest balloon in (1,100),free memory will less than 50M
@@ -70,7 +73,7 @@ class BallooningTest(MemoryBaseTest):
             session = self.vm.wait_for_login(timeout=timeout)
             try:
                 if self.get_win_mon_free_mem(session) > 50:
-                    error.TestFail("Balloon_min test failed %s" % step)
+                    self.test.fail("Balloon_min test failed %s" % step)
             finally:
                 session.close()
         else:
@@ -78,10 +81,10 @@ class BallooningTest(MemoryBaseTest):
             if (abs(mmem - self.ori_mem) != ballooned_mem or
                     (abs(guest_ballooned_mem - ballooned_mem) > 100)):
                 self.error_report(step, self.ori_mem - ballooned_mem, mmem, gmem)
-                raise error.TestFail("Balloon test failed %s" % step)
+                raise exceptions.TestFail("Balloon test failed %s" % step)
         return (mmem, gmem)
 
-    @error.context_aware
+    @error_context.context_aware
     def balloon_memory(self, new_mem):
         """
         Baloon memory to new_mem and verifies on both qemu monitor and
@@ -91,13 +94,14 @@ class BallooningTest(MemoryBaseTest):
         :type new_mem: int
         """
         self.env["balloon_test"] = 0
-        error.context("Change VM memory to %s" % new_mem, logging.info)
+        error_context.context("Change VM memory to %s" % new_mem, logging.info)
         try:
             self.vm.balloon(new_mem)
             self.env["balloon_test"] = 1
         except Exception, e:
             if self.params.get('illegal_value_check', 'no') == 'no' and new_mem != self.get_ballooned_memory():
-                raise error.TestFail("Balloon memory fail with error message: %s" % e)
+                raise exceptions.TestFail("Balloon memory fail with error"
+                                          " message: %s" % e)
         if new_mem > self.ori_mem:
             compare_mem = self.ori_mem
         elif new_mem == 0:
@@ -113,8 +117,8 @@ class BallooningTest(MemoryBaseTest):
                                       self.get_ballooned_memory()),
                                      balloon_timeout)
         if status is None:
-            raise error.TestFail("Failed to balloon memory to expect"
-                                 " value during %ss" % balloon_timeout)
+            raise exceptions.TestFail("Failed to balloon memory to expect"
+                                      " value during %ss" % balloon_timeout)
 
     def run_balloon_sub_test(self, test, params, env, test_tag):
         """
@@ -199,7 +203,7 @@ class BallooningTest(MemoryBaseTest):
         min_size = max(used_size, min_size)
         return min_size, max_size
 
-    @error.context_aware
+    @error_context.context_aware
     def run_ballooning_test(self, expect_mem, tag):
         """
         Run a loop of ballooning test
@@ -214,7 +218,7 @@ class BallooningTest(MemoryBaseTest):
         def _memory_check_after_sub_test():
             try:
                 output = self.memory_check("after subtest", ballooned_mem)
-            except error.TestFail:
+            except exceptions.TestFail:
                 return None
             return output
 
@@ -246,11 +250,12 @@ class BallooningTest(MemoryBaseTest):
             ballooned_mem = abs(self.ori_mem - expect_mem)
             msg = "Wait memory balloon back after "
             msg += params_tag['sub_test_after_balloon']
-            mmem, gmem = utils_misc.wait_for(_memory_check_after_sub_test,
-                                             timeout, sleep_before_check, 5, msg)
-
-            self.current_mmem = mmem
-            self.current_gmem = gmem
+            ret = utils_misc.wait_for(_memory_check_after_sub_test,
+                                      timeout, sleep_before_check, 5, msg)
+            if not ret:
+                self.test.fail("After sub test, memory check failed")
+            self.current_mmem = ret[0]
+            self.current_gmem = ret[1]
         return False
 
     def reset_memory(self):
@@ -348,7 +353,7 @@ class BallooningTestWin(BallooningTest):
             free = float(utils_misc.normalize_data_size(free, order_magnitude="M"))
             return int(free)
         else:
-            error.TestFail("Failed to get windows guest free memory")
+            self.test.fail("Failed to get windows guest free memory")
 
 
 class BallooningTestLinux(BallooningTest):
@@ -383,7 +388,7 @@ class BallooningTestLinux(BallooningTest):
         return int(self.get_total_mem())
 
 
-@error.context_aware
+@error_context.context_aware
 def run(test, params, env):
     """
     Check Memory ballooning, use M when compare memory in this script:
@@ -402,7 +407,7 @@ def run(test, params, env):
         balloon_test = BallooningTestLinux(test, params, env)
 
     for tag in params.objects('test_tags'):
-        error.context("Running %s test" % tag, logging.info)
+        error_context.context("Running %s test" % tag, logging.info)
         params_tag = params.object_params(tag)
         if params_tag.get('expect_memory'):
             expect_mem = int(params_tag.get('expect_memory'))