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
5import logging
6import sys
7
8import common
9from autotest_lib.client.bin import utils
10from autotest_lib.client.common_lib import error
11
12
13def 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