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 5from __future__ import absolute_import 6from __future__ import division 7from __future__ import print_function 8 9import logging 10import sys 11 12import common 13from autotest_lib.client.bin import utils 14from autotest_lib.client.common_lib import error 15import six 16 17 18def cleanup_if_fail(): 19 """Decorator to do cleanup if container fails to be set up. 20 """ 21 def deco_cleanup_if_fail(func): 22 """Wrapper for the decorator. 23 24 @param func: Function to be called. 25 """ 26 def func_cleanup_if_fail(*args, **kwargs): 27 """Decorator to do cleanup if container fails to be set up. 28 29 The first argument must be a ContainerBucket object, which can be 30 used to retrieve the container object by name. 31 32 @param func: function to be called. 33 @param args: arguments for function to be called. 34 @param kwargs: keyword arguments for function to be called. 35 """ 36 bucket = args[0] 37 container_id = utils.get_function_arg_value( 38 func, 'container_id', args, kwargs) 39 try: 40 skip_cleanup = utils.get_function_arg_value( 41 func, 'skip_cleanup', args, kwargs) 42 except (KeyError, ValueError): 43 skip_cleanup = False 44 try: 45 return func(*args, **kwargs) 46 except: 47 exc_info = sys.exc_info() 48 try: 49 container = bucket.get_container(container_id) 50 if container and not skip_cleanup: 51 container.destroy() 52 except error.CmdError as e: 53 logging.error(e) 54 55 # Raise the cached exception with original backtrace. 56 six.reraise(exc_info[0], exc_info[1], exc_info[2]) 57 58 return func_cleanup_if_fail 59 return deco_cleanup_if_fail 60