[telemetry] Retry unit tests after failures in setUpClass and setUpModule.

tools/telemetry/telemetry/unittest/json_results.py

Index: tools/telemetry/telemetry/unittest/json_results.py
diff --git a/tools/telemetry/telemetry/unittest/json_results.py b/tools/telemetry/telemetry/unittest/json_results.py
index 80d7cbe128401313c369a92c4b814472f7d37f9c..256e5c93efecd3d9773c6ce337a606abe3bb0357 100644
--- a/tools/telemetry/telemetry/unittest/json_results.py
+++ b/tools/telemetry/telemetry/unittest/json_results.py
@@ -2,7 +2,9 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import functools
 import json
+import re
 import time
 import unittest
 import urllib2
@@ -87,7 +89,8 @@ def FullResults(args, suite, results):
 
   all_test_names = AllTestNames(suite)
   sets_of_passing_test_names = map(PassingTestNames, results)
-  sets_of_failing_test_names = map(FailedTestNames, results)
+  sets_of_failing_test_names = map(functools.partial(FailedTestNames, suite),
+                                   results)
 
   # TODO(crbug.com/405379): This handles tests that are skipped via the
   # unittest skip decorators (like skipUnless). The tests that are skipped via
@@ -97,7 +100,7 @@ def FullResults(args, suite, results):
                                        - sets_of_failing_test_names[0])
 
   num_tests = len(all_test_names)
-  num_failures = NumFailuresAfterRetries(results)
+  num_failures = NumFailuresAfterRetries(suite, results)
   num_skips = len(skipped_tests)
   num_passes = num_tests - num_failures - num_skips
   full_results['num_failures_by_type'] = {
@@ -160,12 +163,36 @@ def AllTestNames(suite):
   return test_names
 
 
-def NumFailuresAfterRetries(results):
-  return len(FailedTestNames(results[-1]))
+def NumFailuresAfterRetries(suite, results):
+  return len(FailedTestNames(suite, results[-1]))
 
 
-def FailedTestNames(result):
-  return set(test.id() for test, _ in result.failures + result.errors)
+def FailedTestNames(suite, result):
+  failed_test_names = set()
+  for test, error in result.failures + result.errors:
+    if isinstance(test, unittest.TestCase):
+      failed_test_names.add(test.id())
+    elif isinstance(test, unittest.suite._ErrorHolder):  # pylint: disable=W0212
+      # If there's an error in setUpClass or setUpModule, unittest gives us an
+      # _ErrorHolder object. We can parse the object's id for the class or
+      # module that failed, then find all tests in that class or module.
+      match = re.match('setUp[a-zA-Z]+ \\((.+)\\)', test.id())
+      assert match, "Don't know how to retry after this error:\n%s" % error
+      module_or_class = match.groups()[0]
+      failed_test_names |= _FindChildren(module_or_class, AllTestNames(suite))
+    else:
+      assert False, 'Unknown test type: %s' % test.__class__
+  return failed_test_names
+
+
+def _FindChildren(parent, potential_children):
+  children = set()
+  parent_name_parts = parent.split('.')
+  for potential_child in potential_children:
+    child_name_parts = potential_child.split('.')
+    if parent_name_parts == child_name_parts[:len(parent_name_parts)]:
+      children.add(potential_child)
+  return children
 
 
 def PassingTestNames(result):