Handle Telemetry's TimeoutException in retries of flaky tests.

content/test/gpu/gpu_tests/gpu_test_base_unittest.py

Index: content/test/gpu/gpu_tests/gpu_test_base_unittest.py
diff --git a/content/test/gpu/gpu_tests/gpu_test_base_unittest.py b/content/test/gpu/gpu_tests/gpu_test_base_unittest.py
index 37d615493d76dfb83ddb85affcb1e45bb13c9e3a..a80b1d8fdabefac9d0b81a48b3c38bc7d7d4febd 100644
--- a/content/test/gpu/gpu_tests/gpu_test_base_unittest.py
+++ b/content/test/gpu/gpu_tests/gpu_test_base_unittest.py
@@ -45,12 +45,13 @@ class ValidatorWhichFailsNTimes(FakeValidator):
 
 
 class FakePage(gpu_test_base.PageBase):
-  def __init__(self, benchmark, name, manager_mock=None):
+  def __init__(self, benchmark, name, manager_mock=None,
+               shared_page_state_class=gpu_test_base.FakeGpuSharedPageState):
     super(FakePage, self).__init__(
       name=name,
       url='http://nonexistentserver.com/' + name,
       page_set=benchmark.GetFakeStorySet(),
-      shared_page_state_class=gpu_test_base.FakeGpuSharedPageState,
+      shared_page_state_class=shared_page_state_class,
       expectations=benchmark.GetExpectations())
     if manager_mock == None:
       self.RunNavigateSteps = mock.Mock()
@@ -60,6 +61,21 @@ class FakePage(gpu_test_base.PageBase):
       self.RunPageInteractions = manager_mock.RunPageInteractions
 
 
+class SharedPageStateWhichTimesOutNTimes(gpu_test_base.FakeGpuSharedPageState):
+  # This has to be a class variable because it's not possible to
+  # interpose on SharedPageState's constructor.
+  times_to_fail = 1
+
+  def DidRunStory(self, results):  #pylint: disable=arguments-differ
+    if SharedPageStateWhichTimesOutNTimes.times_to_fail > 0:
+      SharedPageStateWhichTimesOutNTimes.times_to_fail = \
+          SharedPageStateWhichTimesOutNTimes.times_to_fail - 1
+      super(SharedPageStateWhichTimesOutNTimes, self).DidRunStory(
+        results, force_timeout_exception=True)
+    super(SharedPageStateWhichTimesOutNTimes, self).DidRunStory(
+      self, results)
+
+
 class FakeTest(gpu_test_base.TestBase):
   def __init__(self,
                times_to_fail_test=0,
@@ -117,9 +133,11 @@ class CrashingPage(FakePage):
 
 
 class PageWhichFailsNTimes(FakePage):
-  def __init__(self, benchmark, name, times_to_fail, manager_mock=None):
-    super(PageWhichFailsNTimes, self).__init__(benchmark, name,
-                                               manager_mock=manager_mock)
+  def __init__(self, benchmark, name, times_to_fail, manager_mock=None,
+               shared_page_state_class=gpu_test_base.FakeGpuSharedPageState):
+    super(PageWhichFailsNTimes, self).__init__(
+      benchmark, name, manager_mock=manager_mock,
+      shared_page_state_class=shared_page_state_class)
     self._times_to_fail = times_to_fail
     self.RunNavigateSteps.side_effect = self.maybeFail
 
@@ -318,6 +336,30 @@ class PageExecutionTest(unittest.TestCase):
                   page, mock.ANY, mock.ANY)]
     self.assertTrue(manager.mock_calls == expected)
 
+  def testFlakyFailureInDidRunStory(self):
+    manager = mock.Mock()
+    test, finder_options = self.setupTest(manager_mock=manager)
+    page = PageWhichFailsNTimes(
+      test, 'page1', 1, manager_mock=manager.page,
+      shared_page_state_class=SharedPageStateWhichTimesOutNTimes)
+    SharedPageStateWhichTimesOutNTimes.times_to_fail = 1
+    test.AddFakePage(page)
+    test.GetExpectations().Flaky('page1')
+    self.assertEqual(test.Run(finder_options), 0,
+                     'Test should run with no errors')
+    expected = [mock.call.validator.WillNavigateToPage(
+                  page, mock.ANY),
+                mock.call.page.RunNavigateSteps(mock.ANY),
+                mock.call.validator.WillNavigateToPage(
+                  page, mock.ANY),
+                mock.call.page.RunNavigateSteps(mock.ANY),
+                mock.call.validator.DidNavigateToPage(
+                  page, mock.ANY),
+                mock.call.page.RunPageInteractions(mock.ANY),
+                mock.call.validator.ValidateAndMeasurePage(
+                  page, mock.ANY, mock.ANY)]
+    self.assertTrue(manager.mock_calls == expected)
+
   def testFlakyPageExceedingNumRetries(self):
     manager = mock.Mock()
     test, finder_options = self.setupTest(manager_mock=manager)