Merge remote branch 'cros/upstream' into master

client/common_lib/error.py

Index: client/common_lib/error.py
diff --git a/client/common_lib/error.py b/client/common_lib/error.py
index 42dfe2bca3d78266b4538f9d9a11d81f00a8d331..0c5641cd8b8cef51696abe655892bf2b22bebef3 100644
--- a/client/common_lib/error.py
+++ b/client/common_lib/error.py
@@ -2,13 +2,14 @@
 Internal global error types
 """
 
-import sys, traceback
+import sys, traceback, threading, logging
 from traceback import format_exception
 
 # Add names you want to be imported by 'from errors import *' to this list.
 # This must be list not a tuple as we modify it to include all of our
 # the Exception classes we define below at the end of this file.
-__all__ = ['format_error']
+__all__ = ['format_error', 'context_aware', 'context', 'get_context',
+           'exception_context']
 
 
 def format_error():
@@ -21,6 +22,141 @@ def format_error():
     return ''.join(trace)
 
 
+# Exception context information:
+# ------------------------------
+# Every function can have some context string associated with it.
+# The context string can be changed by calling context(str) and cleared by
+# calling context() with no parameters.
+# get_context() joins the current context strings of all functions in the
+# provided traceback.  The result is a brief description of what the test was
+# doing in the provided traceback (which should be the traceback of a caught
+# exception).
+#
+# For example: assume a() calls b() and b() calls c().
+#
+# @error.context_aware
+# def a():
+#     error.context("hello")
+#     b()
+#     error.context("world")
+#     error.get_context() ----> 'world'
+#
+# @error.context_aware
+# def b():
+#     error.context("foo")
+#     c()
+#
+# @error.context_aware
+# def c():
+#     error.context("bar")
+#     error.get_context() ----> 'hello --> foo --> bar'
+#
+# The current context is automatically inserted into exceptions raised in
+# context_aware functions, so usually test code doesn't need to call
+# error.get_context().
+
+ctx = threading.local()
+
+
+def _new_context(s=""):
+    if not hasattr(ctx, "contexts"):
+        ctx.contexts = []
+    ctx.contexts.append(s)
+
+
+def _pop_context():
+    ctx.contexts.pop()
+
+
+def context(s="", log=None):
+    """
+    Set the context for the currently executing function and optionally log it.
+
+    @param s: A string.  If not provided, the context for the current function
+            will be cleared.
+    @param log: A logging function to pass the context message to.  If None, no
+            function will be called.
+    """
+    ctx.contexts[-1] = s
+    if s and log:
+        log("Context: %s" % get_context())
+
+
+def base_context(s="", log=None):
+    """
+    Set the base context for the currently executing function and optionally
+    log it.  The base context is just another context level that is hidden by
+    default.  Functions that require a single context level should not use
+    base_context().
+
+    @param s: A string.  If not provided, the base context for the current
+            function will be cleared.
+    @param log: A logging function to pass the context message to.  If None, no
+            function will be called.
+    """
+    ctx.contexts[-1] = ""
+    ctx.contexts[-2] = s
+    if s and log:
+        log("Context: %s" % get_context())
+
+
+def get_context():
+    """Return the current context (or None if none is defined)."""
+    if hasattr(ctx, "contexts"):
+        return " --> ".join([s for s in ctx.contexts if s])
+
+
+def exception_context(e):
+    """Return the context of a given exception (or None if none is defined)."""
+    if hasattr(e, "_context"):
+        return e._context
+
+
+def set_exception_context(e, s):
+    """Set the context of a given exception."""
+    e._context = s
+
+
+def join_contexts(s1, s2):
+    """Join two context strings."""
+    if s1:
+        if s2:
+            return "%s --> %s" % (s1, s2)
+        else:
+            return s1
+    else:
+        return s2
+
+
+def context_aware(fn):
+    """A decorator that must be applied to functions that call context()."""
+    def new_fn(*args, **kwargs):
+        _new_context()
+        _new_context("(%s)" % fn.__name__)
+        try:
+            try:
+                return fn(*args, **kwargs)
+            except Exception, e:
+                if not exception_context(e):
+                    set_exception_context(e, get_context())
+                raise
+        finally:
+            _pop_context()
+            _pop_context()
+    new_fn.__name__ = fn.__name__
+    new_fn.__doc__ = fn.__doc__
+    new_fn.__dict__.update(fn.__dict__)
+    return new_fn
+
+
+def _context_message(e):
+    s = exception_context(e)
+    if s:
+        return "    [context: %s]" % s
+    else:
+        return ""
+
+
 class JobContinue(SystemExit):
     """Allow us to bail out requesting continuance."""
     pass
@@ -33,7 +169,8 @@ class JobComplete(SystemExit):
 
 class AutotestError(Exception):
     """The parent of all errors deliberatly thrown within the client code."""
-    pass
+    def __str__(self):
+        return Exception.__str__(self) + _context_message(self)
 
 
 class JobError(AutotestError):
@@ -46,10 +183,14 @@ class UnhandledJobError(JobError):
     def __init__(self, unhandled_exception):
         if isinstance(unhandled_exception, JobError):
             JobError.__init__(self, *unhandled_exception.args)
+        elif isinstance(unhandled_exception, str):
+            JobError.__init__(self, unhandled_exception)
         else:
             msg = "Unhandled %s: %s"
             msg %= (unhandled_exception.__class__.__name__,
                     unhandled_exception)
+            if not isinstance(unhandled_exception, AutotestError):
+                msg += _context_message(unhandled_exception)
             msg += "\n" + traceback.format_exc()
             JobError.__init__(self, msg)
 
@@ -87,10 +228,14 @@ class UnhandledTestError(TestError):
     def __init__(self, unhandled_exception):
         if isinstance(unhandled_exception, TestError):
             TestError.__init__(self, *unhandled_exception.args)
+        elif isinstance(unhandled_exception, str):
+            TestError.__init__(self, unhandled_exception)
         else:
             msg = "Unhandled %s: %s"
             msg %= (unhandled_exception.__class__.__name__,
                     unhandled_exception)
+            if not isinstance(unhandled_exception, AutotestError):
+                msg += _context_message(unhandled_exception)
             msg += "\n" + traceback.format_exc()
             TestError.__init__(self, msg)
 
@@ -100,10 +245,14 @@ class UnhandledTestFail(TestFail):
     def __init__(self, unhandled_exception):
         if isinstance(unhandled_exception, TestFail):
             TestFail.__init__(self, *unhandled_exception.args)
+        elif isinstance(unhandled_exception, str):
+            TestFail.__init__(self, unhandled_exception)
         else:
             msg = "Unhandled %s: %s"
             msg %= (unhandled_exception.__class__.__name__,
                     unhandled_exception)
+            if not isinstance(unhandled_exception, AutotestError):
+                msg += _context_message(unhandled_exception)
             msg += "\n" + traceback.format_exc()
             TestFail.__init__(self, msg)
 
@@ -118,7 +267,6 @@ class CmdError(TestError):
         self.result_obj = result_obj
         self.additional_text = additional_text
 
-
     def __str__(self):
         if self.result_obj.exit_status is None:
             msg = "Command <%s> failed and is not responding to signals"
@@ -129,6 +277,7 @@ class CmdError(TestError):
 
         if self.additional_text:
             msg += ", " + self.additional_text
+        msg += _context_message(self)
         msg += '\n' + repr(self.result_obj)
         return msg