Fix for flakiness in pyauto automation hook WaitForDownloadsToComplete.

chrome/browser/automation/automation_provider_observers.cc

Index: chrome/browser/automation/automation_provider_observers.cc
diff --git a/chrome/browser/automation/automation_provider_observers.cc b/chrome/browser/automation/automation_provider_observers.cc
index 61dfc6a133bd5aff032ecf22682a36b9298f7b9e..8fd9e5bdc79b6c91fe66840bc988151d39e76d0c 100644
--- a/chrome/browser/automation/automation_provider_observers.cc
+++ b/chrome/browser/automation/automation_provider_observers.cc
@@ -1408,59 +1408,6 @@ void AutomationProviderBookmarkModelObserver::ReplyAndDelete(bool success) {
   delete this;
 }
 
-AutomationProviderDownloadItemObserver::AutomationProviderDownloadItemObserver(
-    AutomationProvider* provider,
-    IPC::Message* reply_message,
-    int downloads)
-    : provider_(provider->AsWeakPtr()),
-      reply_message_(reply_message),
-      downloads_(downloads),
-      interrupted_(false) {
-}
-
-AutomationProviderDownloadItemObserver::
-    ~AutomationProviderDownloadItemObserver() {}
-
-void AutomationProviderDownloadItemObserver::OnDownloadUpdated(
-    DownloadItem* download) {
-  interrupted_ |= download->IsInterrupted();
-  // If any download was interrupted, on the next update each outstanding
-  // download is cancelled.
-  if (interrupted_) {
-    // |Cancel()| does nothing if |download| is already interrupted.
-    download->Cancel(true);
-    RemoveAndCleanupOnLastEntry(download);
-  }
-
-  if (download->IsComplete())
-    RemoveAndCleanupOnLastEntry(download);
-}
-
-// We don't want to send multiple messages, as the behavior is undefined.
-// Set |interrupted_| on error, and on the last download completed/
-// interrupted, send either an error or a success message.
-void AutomationProviderDownloadItemObserver::RemoveAndCleanupOnLastEntry(
-    DownloadItem* download) {
-  // Forget about the download.
-  download->RemoveObserver(this);
-  if (--downloads_ == 0) {
-    if (provider_) {
-      if (interrupted_) {
-        AutomationJSONReply(provider_, reply_message_.release()).SendError(
-            "Download Interrupted");
-      } else {
-        AutomationJSONReply(provider_, reply_message_.release()).SendSuccess(
-            NULL);
-      }
-    }
-    delete this;
-  }
-}
-
-void AutomationProviderDownloadItemObserver::OnDownloadOpened(
-    DownloadItem* download) {
-}
-
 AutomationProviderDownloadUpdatedObserver::
 AutomationProviderDownloadUpdatedObserver(
     AutomationProvider* provider,
@@ -1526,6 +1473,61 @@ void AutomationProviderDownloadModelChangedObserver::ModelChanged() {
   delete this;
 }
 
+AllDownloadsCompleteObserver::AllDownloadsCompleteObserver(
+    AutomationProvider* provider,
+    IPC::Message* reply_message,
+    DownloadManager* download_manager,
+    ListValue* pre_download_ids)
+    : provider_(provider->AsWeakPtr()),
+      reply_message_(reply_message),
+      download_manager_(download_manager) {
+  for (ListValue::iterator it = pre_download_ids->begin();
+       it != pre_download_ids->end(); ++it) {
+    int val = 0;
+    (*it)->GetAsInteger(&val);

[Nirnimesh, 2011/08/08 17:26:17]

wrap this inside if()

[dennis_jeffrey, 2011/08/08 22:29:51]

Done.

+    pre_download_ids_.insert(val);
+  }
+  download_manager_->AddObserver(this);  // Will call initial ModelChanged().
+}
+
+AllDownloadsCompleteObserver::~AllDownloadsCompleteObserver() {}
+
+void AllDownloadsCompleteObserver::ModelChanged() {
+  // The set of downloads in the download manager has changed.  If there are
+  // any new downloads that are still in progress, add them to the pending list.
+  std::vector<DownloadItem*> downloads;
+  download_manager_->GetAllDownloads(FilePath(), &downloads);
+  for (std::vector<DownloadItem*>::iterator it = downloads.begin();
+       it != downloads.end(); ++it) {
+    if ((*it)->state() == DownloadItem::IN_PROGRESS &&
+        pre_download_ids_.find((*it)->id()) == pre_download_ids_.end()) {
+      (*it)->AddObserver(this);
+      pending_downloads_.insert(*it);
+    }
+  }
+  ReplyIfNecessary();
+}
+
+void AllDownloadsCompleteObserver::OnDownloadUpdated(DownloadItem* download) {
+  // If the current download's status has changed to a final state (not state
+  // "in progress"), remove it from the pending list.
+  if (download->state() != DownloadItem::IN_PROGRESS) {

[Nirnimesh, 2011/08/08 17:26:17]

shouldn't you check for status COMPLETED?

[dennis_jeffrey, 2011/08/08 22:29:51]

No, I want the hook to return as soon as we don't need to wait for any downloads
to further progress.  The "IN_PROGRESS" state is the only state in which a
download is still actively progressing.  The other states -- COMPLETE,
CANCELLED, REMOVING, INTERRUPTED -- occur when a download will not make any
further progression.  I'm leaving it up to the test to check/decide what to do
with the end result of a download; this hook merely waits until all downloads
stop progressing.

If we only wait for the COMPLETE state, this hook may never return.

+    download->RemoveObserver(this);
+    pending_downloads_.erase(download);
+    ReplyIfNecessary();
+  }
+}
+
+void AllDownloadsCompleteObserver::ReplyIfNecessary() {
+  if (pending_downloads_.empty()) {

[Paweł Hajdan Jr., 2011/08/08 16:36:13]

nit: Preferred would be:

if (!pending_downloads_.empty())
  return;

[dennis_jeffrey, 2011/08/08 22:29:51]

Done.

+    download_manager_->RemoveObserver(this);
+    if (provider_)

[Paweł Hajdan Jr., 2011/08/08 16:36:13]

nit: Add braces {} here.

[dennis_jeffrey, 2011/08/08 22:29:51]

I assume this is to improve readability because lines 1525-1526 below is on 2
lines?  If so, then with the latest changes, those 2 lines can now fit on 1
line, so I think it's ok to leave out the {} now.

+      AutomationJSONReply(provider_,
+                          reply_message_.release()).SendSuccess(NULL);
+    delete this;
+  }
+}
+
 AutomationProviderSearchEngineObserver::AutomationProviderSearchEngineObserver(
     AutomationProvider* provider,
     IPC::Message* reply_message)