Make WaitForInfobarCount() robust

chrome/test/pyautolib/pyauto.py

 #!/usr/bin/python
 
 # Copyright (c) 2010 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """PyAuto: Python Interface to Chromium's Automation Proxy.
 
 PyAuto uses swig to expose Automation Proxy interfaces to Python.
 For complete documentation on the functionality available,
@@ skipping 230 matching lines @@
     return ret
 
   @staticmethod
   def Kill(pid):
     """Terminate the given pid."""
     if PyUITest.IsWin():
       subprocess.call(['taskkill.exe', '/T', '/F', '/PID', str(pid)])
     else:
       os.kill(pid, signal.SIGTERM)
 
-  def WaitUntil(self, function, timeout=-1, retry_sleep=0.25, args=[]):
+  def WaitUntil(self, function, timeout=-1, retry_sleep=0.25, args=[],
+                expect_retval=None):
     """Poll on a condition until timeout.
 
-    Waits until the |function| evalues to True or until |timeout| secs,
-    whichever occurs earlier.
+    Waits until the |function| evalues to |expect_retval| or until |timeout|
+    secs, whichever occurs earlier.
 
     This is better than using a sleep, since it waits (almost) only as much
     as needed.
 
     WARNING: This method call should be avoided as far as possible in favor
     of a real wait from chromium (like wait-until-page-loaded).
     Only use in case there's really no better option.
 
     EXAMPLES:-
     Wait for "file.txt" to get created:
       WaitUntil(os.path.exists, args=["file.txt"])
 
     Same as above, but using lambda:
       WaitUntil(lambda: os.path.exists("file.txt"))
 
     Args:
       function: the function whose truth value is to be evaluated
       timeout: the max timeout (in secs) for which to wait. The default
                action is to wait for kWaitForActionMaxMsec, as set in
                ui_test.cc
                Use None to wait indefinitely.
       retry_sleep: the sleep interval (in secs) before retrying |function|.
                    Defaults to 0.25 secs.
       args: the args to pass to |function|
+      expect_retval: the expected return value for |function|. This forms the
+                     exit criteria. In case this is None (the default),
+                     |function|'s return value is checked for truth,
+                     so 'non-empty-string' should match with True
 
     Returns:
       True, if returning when |function| evaluated to True
       False, when returning due to timeout
     """
     if timeout == -1:  # Default
-      timeout = self.action_max_timeout_ms()/1000.0
+      timeout = self.action_max_timeout_ms() / 1000.0
     assert callable(function), "function should be a callable"
     begin = time.time()
     while timeout is None or time.time() - begin <= timeout:
-      if function(*args):
+      retval = function(*args)
+      if (expect_retval is None and retval) or expect_retval == retval:
         return True
+      logging.debug('WaitUntil(%s) still waiting. '
+                    'Expecting %s. Last returned %s.' % (
+                    function, expect_retval, retval))
       time.sleep(retry_sleep)
     return False
 
   class CmdExecutionTimeoutChanger(object):
     """Facilitate temporary changes to command_execution_timeout_ms.
 
     Automatically resets to original timeout when object is destroyed.
     """
     _saved_timeout = -1  # Saved value for command_execution_timeout_ms
 
@@ skipping 404 matching lines @@
     Note: Wait duration is capped by the automation timeout.
 
     Args:
       count: requested number of infobars
       windex: window index.  Defaults to 0 (first window)
       tab_index: tab index  Defaults to 0 (first tab)
 
     Raises:
       pyauto_errors.JSONInterfaceError if the automation call returns an error.
     """
-    cmd_dict = {
-      'command': 'WaitForInfobarCount',
-      'count': count,
-      'tab_index': tab_index,
-    }
     # TODO(phajdan.jr): We need a solid automation infrastructure to handle
     # these cases. See crbug.com/53647.
-    return self.WaitUntil(
-        lambda(count): len(self.GetBrowserInfo()\
-            ['windows'][windex]['tabs'][tab_index]['infobars']) == count,
-        args=[count])
+    def _InfobarCount():
+      windows = self.GetBrowserInfo()['windows']
+      if windex >= len(windows):  # not enough windows
+        return -1
+      tabs = windows[windex]['tabs']
+      if tab_index >= len(tabs):  # not enough tabs
+        return -1
+      return len(tabs[tab_index]['infobars'])
+
+    return self.WaitUntil(_InfobarCount, expect_retval=count)
 
   def PerformActionOnInfobar(
       self, action, infobar_index, windex=0, tab_index=0):
     """Perform actions on an infobar.
 
     Args:
       action: the action to be performed.
               Actions depend on the type of the infobar.  The user needs to
               call the right action for the right infobar.
               Valid inputs are:
@@ skipping 1055 matching lines @@
     if self._options.verbose:
       verbosity = 2
     result = PyAutoTextTestRuner(verbosity=verbosity).run(pyauto_suite)
     del loaded_tests  # Need to destroy test cases before the suite
     del pyauto_suite
     sys.exit(not result.wasSuccessful())
 
 
 if __name__ == '__main__':
   Main()