Fix for flakiness in pyauto automation hook WaitForDownloadsToComplete.

chrome/browser/automation/automation_provider_observers.h

 // 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.
 
 #ifndef CHROME_BROWSER_AUTOMATION_AUTOMATION_PROVIDER_OBSERVERS_H_
 #define CHROME_BROWSER_AUTOMATION_AUTOMATION_PROVIDER_OBSERVERS_H_
 #pragma once
 
 #include <deque>
 #include <map>
 #include <set>
 #include <string>
 #include <vector>
 
 #include "base/compiler_specific.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/memory/weak_ptr.h"
+#include "base/values.h"
 #include "base/synchronization/waitable_event.h"
 #include "chrome/browser/automation/automation_provider_json.h"
 #include "chrome/browser/automation/automation_tab_helper.h"
 #include "chrome/browser/bookmarks/bookmark_model_observer.h"
 #include "chrome/browser/browsing_data_remover.h"
 #if defined(OS_CHROMEOS)
 #include "chrome/browser/chromeos/cros/network_library.h"
 #include "chrome/browser/chromeos/login/enterprise_enrollment_screen.h"
 #include "chrome/browser/chromeos/login/enterprise_enrollment_screen_actor.h"
 #include "chrome/browser/chromeos/login/enterprise_enrollment_view.h"
@@ skipping 1000 matching lines @@
   // observer list).
   void ReplyAndDelete(bool success);
 
   base::WeakPtr<AutomationProvider> automation_provider_;
   scoped_ptr<IPC::Message> reply_message_;
   BookmarkModel* model_;
 
   DISALLOW_COPY_AND_ASSIGN(AutomationProviderBookmarkModelObserver);
 };
 
-// Allows the automation provider to wait for all downloads to finish.
-// If any download is interrupted, it will cancel all the other downloads at
-// the next |OnDownloadUpdated|, and send an error when all are done.
-class AutomationProviderDownloadItemObserver : public DownloadItem::Observer {
- public:
-  AutomationProviderDownloadItemObserver(
-      AutomationProvider* provider,
-      IPC::Message* reply_message,
-      int downloads);
-  virtual ~AutomationProviderDownloadItemObserver();
-
-  virtual void OnDownloadUpdated(DownloadItem* download);
-  virtual void OnDownloadOpened(DownloadItem* download);
-
- private:
-  void RemoveAndCleanupOnLastEntry(DownloadItem* download);
-
-  base::WeakPtr<AutomationProvider> provider_;
-  scoped_ptr<IPC::Message> reply_message_;
-  int downloads_;
-  bool interrupted_;
-
-  DISALLOW_COPY_AND_ASSIGN(AutomationProviderDownloadItemObserver);
-};
-
 // Allows the automation provider to wait until the download has been updated
 // or opened.
 class AutomationProviderDownloadUpdatedObserver
     : public DownloadItem::Observer {
  public:
   AutomationProviderDownloadUpdatedObserver(
       AutomationProvider* provider,
       IPC::Message* reply_message,
       bool wait_for_open);
   virtual ~AutomationProviderDownloadUpdatedObserver();
@@ skipping 23 matching lines @@
   virtual void ModelChanged();
 
  private:
   base::WeakPtr<AutomationProvider> provider_;
   scoped_ptr<IPC::Message> reply_message_;
   DownloadManager* download_manager_;
 
   DISALLOW_COPY_AND_ASSIGN(AutomationProviderDownloadModelChangedObserver);
 };
 
+// Observes when all pending downloads have completed.
+class AllDownloadsCompleteObserver
+    : public DownloadManager::Observer,
+      public DownloadItem::Observer {
+ public:
+  AllDownloadsCompleteObserver(
+      AutomationProvider* provider,
+      IPC::Message* reply_message,
+      DownloadManager* download_manager,
+      ListValue* pre_download_ids);
+  virtual ~AllDownloadsCompleteObserver();
+
+  // DownloadManager::Observer.
+  virtual void ModelChanged();
+
+  // DownloadItem::Observer.
+  virtual void OnDownloadUpdated(DownloadItem* download);
+  virtual void OnDownloadOpened(DownloadItem* download) {}
+
+ private:
+  void ReplyIfNecessary();
+
+  base::WeakPtr<AutomationProvider> provider_;
+  scoped_ptr<IPC::Message> reply_message_;
+  DownloadManager* download_manager_;
+  std::set<int> pre_download_ids_;
+  std::set<DownloadItem*> pending_downloads_;
+
+  DISALLOW_COPY_AND_ASSIGN(AllDownloadsCompleteObserver);
+};
+
 // Allows automation provider to wait until TemplateURLService has loaded
 // before looking up/returning search engine info.
 class AutomationProviderSearchEngineObserver
     : public TemplateURLServiceObserver {
  public:
   AutomationProviderSearchEngineObserver(
       AutomationProvider* provider,
       IPC::Message* reply_message);
   virtual ~AutomationProviderSearchEngineObserver();
 
@@ skipping 530 matching lines @@
 
  private:
   NotificationRegistrar registrar_;
   base::WeakPtr<AutomationProvider> automation_;
   scoped_ptr<IPC::Message> reply_message_;
 
   DISALLOW_COPY_AND_ASSIGN(DragTargetDropAckNotificationObserver);
 };
 
 #endif  // CHROME_BROWSER_AUTOMATION_AUTOMATION_PROVIDER_OBSERVERS_H_