1# Copyright 2014 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"""
6Function/method decorators that provide timeout and retry logic.
7"""
8
9import functools
10import os
11import sys
12
13from pylib import constants
14from pylib.device import device_errors
15from pylib.utils import reraiser_thread
16from pylib.utils import timeout_retry
17
18# TODO(jbudorick) Remove once the DeviceUtils implementations are no longer
19#                 backed by AndroidCommands / android_testrunner.
20sys.path.append(os.path.join(constants.DIR_SOURCE_ROOT, 'third_party',
21                             'android_testrunner'))
22import errors as old_errors
23
24DEFAULT_TIMEOUT_ATTR = '_default_timeout'
25DEFAULT_RETRIES_ATTR = '_default_retries'
26
27
28def _TimeoutRetryWrapper(f, timeout_func, retries_func, pass_values=False):
29  """ Wraps a funcion with timeout and retry handling logic.
30
31  Args:
32    f: The function to wrap.
33    timeout_func: A callable that returns the timeout value.
34    retries_func: A callable that returns the retries value.
35    pass_values: If True, passes the values returned by |timeout_func| and
36                 |retries_func| to the wrapped function as 'timeout' and
37                 'retries' kwargs, respectively.
38  Returns:
39    The wrapped function.
40  """
41  @functools.wraps(f)
42  def TimeoutRetryWrapper(*args, **kwargs):
43    timeout = timeout_func(*args, **kwargs)
44    retries = retries_func(*args, **kwargs)
45    if pass_values:
46      kwargs['timeout'] = timeout
47      kwargs['retries'] = retries
48    def impl():
49      return f(*args, **kwargs)
50    try:
51      return timeout_retry.Run(impl, timeout, retries)
52    except old_errors.WaitForResponseTimedOutError as e:
53      raise device_errors.CommandTimeoutError(str(e)), None, (
54             sys.exc_info()[2])
55    except old_errors.DeviceUnresponsiveError as e:
56      raise device_errors.DeviceUnreachableError(str(e)), None, (
57             sys.exc_info()[2])
58    except reraiser_thread.TimeoutError as e:
59      raise device_errors.CommandTimeoutError(str(e)), None, (
60             sys.exc_info()[2])
61  return TimeoutRetryWrapper
62
63
64def WithTimeoutAndRetries(f):
65  """A decorator that handles timeouts and retries.
66
67  'timeout' and 'retries' kwargs must be passed to the function.
68
69  Args:
70    f: The function to decorate.
71  Returns:
72    The decorated function.
73  """
74  get_timeout = lambda *a, **kw: kw['timeout']
75  get_retries = lambda *a, **kw: kw['retries']
76  return _TimeoutRetryWrapper(f, get_timeout, get_retries)
77
78
79def WithExplicitTimeoutAndRetries(timeout, retries):
80  """Returns a decorator that handles timeouts and retries.
81
82  The provided |timeout| and |retries| values are always used.
83
84  Args:
85    timeout: The number of seconds to wait for the decorated function to
86             return. Always used.
87    retries: The number of times the decorated function should be retried on
88             failure. Always used.
89  Returns:
90    The actual decorator.
91  """
92  def decorator(f):
93    get_timeout = lambda *a, **kw: timeout
94    get_retries = lambda *a, **kw: retries
95    return _TimeoutRetryWrapper(f, get_timeout, get_retries)
96  return decorator
97
98
99def WithTimeoutAndRetriesDefaults(default_timeout, default_retries):
100  """Returns a decorator that handles timeouts and retries.
101
102  The provided |default_timeout| and |default_retries| values are used only
103  if timeout and retries values are not provided.
104
105  Args:
106    default_timeout: The number of seconds to wait for the decorated function
107                     to return. Only used if a 'timeout' kwarg is not passed
108                     to the decorated function.
109    default_retries: The number of times the decorated function should be
110                     retried on failure. Only used if a 'retries' kwarg is not
111                     passed to the decorated function.
112  Returns:
113    The actual decorator.
114  """
115  def decorator(f):
116    get_timeout = lambda *a, **kw: kw.get('timeout', default_timeout)
117    get_retries = lambda *a, **kw: kw.get('retries', default_retries)
118    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)
119  return decorator
120
121
122def WithTimeoutAndRetriesFromInstance(
123    default_timeout_name=DEFAULT_TIMEOUT_ATTR,
124    default_retries_name=DEFAULT_RETRIES_ATTR):
125  """Returns a decorator that handles timeouts and retries.
126
127  The provided |default_timeout_name| and |default_retries_name| are used to
128  get the default timeout value and the default retries value from the object
129  instance if timeout and retries values are not provided.
130
131  Note that this should only be used to decorate methods, not functions.
132
133  Args:
134    default_timeout_name: The name of the default timeout attribute of the
135                          instance.
136    default_retries_name: The name of the default retries attribute of the
137                          instance.
138  Returns:
139    The actual decorator.
140  """
141  def decorator(f):
142    def get_timeout(inst, *_args, **kwargs):
143      return kwargs.get('timeout', getattr(inst, default_timeout_name))
144    def get_retries(inst, *_args, **kwargs):
145      return kwargs.get('retries', getattr(inst, default_retries_name))
146    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)
147  return decorator
148
149