Fix for flakiness in pyauto automation hook WaitForDownloadsToComplete.

chrome/browser/automation/testing_automation_provider.cc

 // Copyright (c) 2011 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.
 
 #include "chrome/browser/automation/testing_automation_provider.h"
 
 #include <map>
 #include <set>
 #include <string>
 #include <vector>
@@ skipping 2314 matching lines @@
   browser_handler_map["GetPrefsInfo"] =
       &TestingAutomationProvider::GetPrefsInfo;
   browser_handler_map["SetPrefs"] = &TestingAutomationProvider::SetPrefs;
 
   browser_handler_map["SetWindowDimensions"] =
       &TestingAutomationProvider::SetWindowDimensions;
 
   browser_handler_map["GetDownloadsInfo"] =
       &TestingAutomationProvider::GetDownloadsInfo;
   browser_handler_map["WaitForAllDownloadsToComplete"] =
-      &TestingAutomationProvider::WaitForDownloadsToComplete;
+      &TestingAutomationProvider::WaitForAllDownloadsToComplete;
   browser_handler_map["PerformActionOnDownload"] =
       &TestingAutomationProvider::PerformActionOnDownload;
 
   browser_handler_map["GetInitialLoadTimes"] =
       &TestingAutomationProvider::GetInitialLoadTimes;
 
   browser_handler_map["SaveTabContents"] =
       &TestingAutomationProvider::SaveTabContents;
 
   browser_handler_map["ImportSettings"] =
@@ skipping 585 matching lines @@
     for (std::vector<DownloadItem*>::iterator it = downloads.begin();
          it != downloads.end();
          it++) {  // Fill info about each download item.
       list_of_downloads->Append(GetDictionaryFromDownloadItem(*it));
     }
   }
   return_value->Set("downloads", list_of_downloads);
   reply.SendSuccess(return_value.get());
 }
 
-void TestingAutomationProvider::WaitForDownloadsToComplete(
+void TestingAutomationProvider::WaitForAllDownloadsToComplete(
     Browser* browser,
     DictionaryValue* args,
     IPC::Message* reply_message) {
+  ListValue* pre_download_ids = NULL;
 
-  // Look for a quick return.
+  if (!args->GetList("pre_download_ids", &pre_download_ids)) {
+    AutomationJSONReply(this, reply_message)
+        .SendError(StringPrintf("List of IDs of previous downloads required."));
+    return;
+  }
   if (!browser->profile()->HasCreatedDownloadManager()) {
-    // No download manager.
+    // No download manager, so no downloads to wait for.
     AutomationJSONReply(this, reply_message).SendSuccess(NULL);
     return;
   }
-  std::vector<DownloadItem*> downloads;
-  browser->profile()->GetDownloadManager()->
-      GetCurrentDownloads(FilePath(), &downloads);
-  if (downloads.empty()) {
-    AutomationJSONReply(this, reply_message).SendSuccess(NULL);
-    return;
-  }
-  // The observer owns itself.  When the last observed item pings, it
-  // deletes itself.
-  AutomationProviderDownloadItemObserver* item_observer =
-      new AutomationProviderDownloadItemObserver(
-          this, reply_message, downloads.size());
-  for (std::vector<DownloadItem*>::iterator i = downloads.begin();
-       i != downloads.end();
-       i++) {
-    (*i)->AddObserver(item_observer);
-  }
+
+  // This observer will delete itself.
+  new AllDownloadsCompleteObserver(
+      this, reply_message, browser->profile()->GetDownloadManager(),
+      pre_download_ids);
 }
 
 namespace {
 
 DownloadItem* GetDownloadItemFromId(int id, DownloadManager* download_manager) {
   std::vector<DownloadItem*> downloads;
   download_manager->GetAllDownloads(FilePath(), &downloads);
   DownloadItem* selected_item = NULL;
 
   for (std::vector<DownloadItem*>::iterator it = downloads.begin();
@@ skipping 3181 matching lines @@
   IPC::ParamTraits<std::vector<GURL> >::Write(reply_message_, redirects_gurl);
 
   Send(reply_message_);
   redirect_query_ = 0;
   reply_message_ = NULL;
 }
 
 void TestingAutomationProvider::OnRemoveProvider() {
   AutomationProviderList::GetInstance()->RemoveProvider(this);
 }