SafeBrowsingSystemTest: use notification for end of update instead of polling.

chrome/browser/safe_browsing/safe_browsing_test.cc

 // Copyright (c) 2012 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.
 //
 // This test uses the safebrowsing test server published at
 // http://code.google.com/p/google-safe-browsing/ to test the safebrowsing
 // protocol implemetation. Details of the safebrowsing testing flow is
 // documented at
 // http://code.google.com/p/google-safe-browsing/wiki/ProtocolTesting
 //
@@ skipping 15 matching lines @@
 #include "base/stringprintf.h"
 #include "base/synchronization/lock.h"
 #include "base/test/test_timeouts.h"
 #include "base/threading/platform_thread.h"
 #include "base/threading/thread.h"
 #include "base/time.h"
 #include "base/utf_string_conversions.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/safe_browsing/protocol_manager.h"
 #include "chrome/browser/safe_browsing/safe_browsing_service.h"
+#include "chrome/common/chrome_notification_types.h"
 #include "chrome/common/chrome_switches.h"
 #include "chrome/common/url_constants.h"
 #include "chrome/test/base/in_process_browser_test.h"
 #include "chrome/test/base/ui_test_utils.h"
 #include "content/public/browser/browser_context.h"
 #include "content/public/test/test_browser_thread.h"
 #include "net/base/host_resolver.h"
 #include "net/base/load_flags.h"
 #include "net/base/net_log.h"
 #include "net/test/python_utils.h"
@@ skipping 174 matching lines @@
 
   void UpdateSafeBrowsingStatus() {
     ASSERT_TRUE(safe_browsing_service_);
     base::AutoLock lock(update_status_mutex_);
     last_update_ = safe_browsing_service_->protocol_manager_->last_update();
     is_update_scheduled_ =
         safe_browsing_service_->protocol_manager_->update_timer_.IsRunning();
   }
 
   void ForceUpdate() {
+    content::WindowedNotificationObserver observer(
+        chrome::NOTIFICATION_SAFE_BROWSING_UPDATE_COMPLETE,
+        content::Source<SafeBrowsingService>(safe_browsing_service_));
+    BrowserThread::PostTask(BrowserThread::IO, FROM_HERE,
+        base::Bind(&SafeBrowsingServiceTest::ForceUpdateOnIOThread,
+                   this));
+    observer.Wait();
+  }
+
+  void ForceUpdateOnIOThread() {
+    EXPECT_TRUE(BrowserThread::CurrentlyOn(BrowserThread::IO));
     ASSERT_TRUE(safe_browsing_service_);
     safe_browsing_service_->protocol_manager_->ForceScheduleNextUpdate(0);
   }
 
   void CheckIsDatabaseReady() {
     base::AutoLock lock(update_status_mutex_);
     is_database_ready_ =
         !safe_browsing_service_->database_update_in_progress_;
   }
 
@@ skipping 131 matching lines @@
   virtual void OnDownloadUrlCheckResult(
       const std::vector<GURL>& url_chain,
       SafeBrowsingService::UrlCheckResult result) {
     // TODO(lzheng): Add test for DownloadUrl.
   }
 
   virtual void OnBlockingPageComplete(bool proceed) {
     NOTREACHED() << "Not implemented.";
   }
 
-  // Functions and callbacks to start the safebrowsing database update.
-  void ForceUpdate() {
-    BrowserThread::PostTask(BrowserThread::IO, FROM_HERE,
-        base::Bind(&SafeBrowsingServiceTestHelper::ForceUpdateInIOThread,
-                   this));
-    // Will continue after OnForceUpdateDone().
-    content::RunMessageLoop();
-  }
-  void ForceUpdateInIOThread() {
-    EXPECT_TRUE(BrowserThread::CurrentlyOn(BrowserThread::IO));
-    safe_browsing_test_->ForceUpdate();
-    BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-        base::Bind(&SafeBrowsingServiceTestHelper::OnForceUpdateDone,
-                   this));
-  }
-  void OnForceUpdateDone() {
-    StopUILoop();
-  }
-
   // Functions and callbacks related to CheckUrl. These are used to verify
   // phishing URLs.
   void CheckUrl(const GURL& url) {
     BrowserThread::PostTask(BrowserThread::IO, FROM_HERE,
         base::Bind(&SafeBrowsingServiceTestHelper::CheckUrlOnIOThread,
                    this, url));
     content::RunMessageLoop();
   }
   void CheckUrlOnIOThread(const GURL& url) {
     EXPECT_TRUE(BrowserThread::CurrentlyOn(BrowserThread::IO));
@@ skipping 27 matching lines @@
     safe_browsing_test_->CheckIsDatabaseReady();
     BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
         base::Bind(&SafeBrowsingServiceTestHelper::OnWaitForStatusUpdateDone,
                    this));
   }
 
   void OnWaitForStatusUpdateDone() {
     StopUILoop();
   }
 
-  // Wait for a given period to get safebrowsing status updated.
-  void WaitForStatusUpdate(base::TimeDelta wait_time) {
-    BrowserThread::PostDelayedTask(
+  // Update safebrowsing status.
+  void UpdateStatus() {
+    BrowserThread::PostTask(
         BrowserThread::IO,
         FROM_HERE,
         base::Bind(&SafeBrowsingServiceTestHelper::CheckStatusOnIOThread,
-                   this),
-        wait_time);
+                   this));
     // Will continue after OnWaitForStatusUpdateDone().
     content::RunMessageLoop();
   }
 
   void WaitTillServerReady(const char* host, int port) {
     response_status_ = net::URLRequestStatus::FAILED;
     GURL url(base::StringPrintf("http://%s:%d%s?test_step=0",
                                 host, port, kDBResetPath));
     // TODO(lzheng): We should have a way to reliably tell when a server is
     // ready so we could get rid of the Sleep and retry loop.
@@ skipping 98 matching lines @@
   SafeBrowsingTestServer test_server(datafile_path);
   ASSERT_TRUE(test_server.Start());
 
   // Make sure the server is running.
   safe_browsing_helper->WaitTillServerReady(server_host, server_port);
 
   // Waits and makes sure safebrowsing update is not happening.
   // The wait will stop once OnWaitForStatusUpdateDone in
   // safe_browsing_helper is called and status from safe_browsing_service_
   // is checked.
-  safe_browsing_helper->WaitForStatusUpdate(base::TimeDelta());
+  safe_browsing_helper->UpdateStatus();
   EXPECT_TRUE(is_database_ready());
   EXPECT_FALSE(is_update_scheduled());
   EXPECT_TRUE(last_update().is_null());
   // Starts updates. After each update, the test will fetch a list of URLs with
   // expected results to verify with safebrowsing service. If there is no error,
   // the test moves on to the next step to get more update chunks.
   // This repeats till there is no update data.
   for (int step = 1;; step++) {
     // Every step should be a fresh start.
     SCOPED_TRACE(base::StringPrintf("step=%d", step));
     EXPECT_TRUE(is_database_ready());
     EXPECT_FALSE(is_update_scheduled());
 
     // Starts safebrowsing update on IO thread. Waits till scheduled
-    // update finishes. Stops waiting after kMaxWaitSecPerStep if the update
-    // could not finish.
+    // update finishes.
     base::Time now = base::Time::Now();
     SetTestStep(step);
-    safe_browsing_helper->ForceUpdate();
+    ForceUpdate();
 
-    // TODO(mattm): use NOTIFICATION_SAFE_BROWSING_UPDATE_COMPLETE instead.
-    do {
-      // Periodically pull the status.
-      safe_browsing_helper->WaitForStatusUpdate(
-          TestTimeouts::tiny_timeout());
-    } while (is_update_scheduled() || !is_database_ready());
-
-
+    safe_browsing_helper->UpdateStatus();
+    EXPECT_TRUE(is_database_ready());
+    EXPECT_FALSE(is_update_scheduled());
+    EXPECT_FALSE(last_update().is_null());
     if (last_update() < now) {
       // This means no data available anymore.
       break;
     }
 
     // Fetches URLs to verify and waits till server responses with data.
     EXPECT_EQ(net::URLRequestStatus::SUCCESS,
               safe_browsing_helper->FetchUrlsToVerify(server_host,
                                                       server_port,
                                                       step));
@@ skipping 33 matching lines @@
   }
 
   // Verifies with server if test is done and waits till server responses.
   EXPECT_EQ(net::URLRequestStatus::SUCCESS,
             safe_browsing_helper->VerifyTestComplete(server_host,
                                                      server_port,
                                                      last_step));
   EXPECT_EQ("yes", safe_browsing_helper->response_data());
   test_server.Stop();
 }