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 container_id = utils.get_function_arg_value( 33 func, 'container_id', args, kwargs) 34 try: 35 skip_cleanup = utils.get_function_arg_value( 36 func, 'skip_cleanup', args, kwargs) 37 except (KeyError, ValueError): 38 skip_cleanup = False 39 try: 40 return func(*args, **kwargs) 41 except: 42 exc_info = sys.exc_info() 43 try: 44 container = bucket.get_container(container_id) 45 if container and not skip_cleanup: 46 container.destroy() 47 except error.CmdError as e: 48 logging.error(e) 49 50 # Raise the cached exception with original backtrace. 51 raise exc_info[0], exc_info[1], exc_info[2] 52 return func_cleanup_if_fail 53 return deco_cleanup_if_fail 54
