1 # Copyright 2015 The Chromium Authors. All rights reserved. 2 # Use of this source code is governed by a BSD-style license that can be 3 # found in the LICENSE file. 4 5 import logging 6 import sys 7 8 import common 9 from autotest_lib.client.bin import utils 10 from autotest_lib.client.common_lib import error 11 12 13 def cleanup_if_fail(): 14 """Decorator to do cleanup if container fails to be set up. 15 """ 16 def deco_cleanup_if_fail(func): 17 """Wrapper for the decorator. 18 19 @param func: Function to be called. 20 """ 21 def func_cleanup_if_fail(*args, **kwargs): 22 """Decorator to do cleanup if container fails to be set up. 23 24 The first argument must be a ContainerBucket object, which can be 25 used to retrieve the container object by name. 26 27 @param func: function to be called. 28 @param args: arguments for function to be called. 29 @param kwargs: keyword arguments for function to be called. 30 """ 31 bucket = args[0] 32 name = utils.get_function_arg_value(func, 'name', args, kwargs) 33 try: 34 skip_cleanup = utils.get_function_arg_value( 35 func, 'skip_cleanup', args, kwargs) 36 except (KeyError, ValueError): 37 skip_cleanup = False 38 try: 39 return func(*args, **kwargs) 40 except: 41 exc_info = sys.exc_info() 42 try: 43 container = bucket.get(name) 44 if container and not skip_cleanup: 45 container.destroy() 46 except error.CmdError as e: 47 logging.error(e) 48 49 # Raise the cached exception with original backtrace. 50 raise exc_info[0], exc_info[1], exc_info[2] 51 return func_cleanup_if_fail 52 return deco_cleanup_if_fail 53
