Create test fixture for SafeBrowsingService

chrome/browser/prerender/prerender_browsertest.cc

Index: chrome/browser/prerender/prerender_browsertest.cc
diff --git a/chrome/browser/prerender/prerender_browsertest.cc b/chrome/browser/prerender/prerender_browsertest.cc
index 9eec60fa294f02a3d36490018784ad8e1cc5b9c2..cafa98ee6b1fb8b1336d502bc49307d57c797fb1 100644
--- a/chrome/browser/prerender/prerender_browsertest.cc
+++ b/chrome/browser/prerender/prerender_browsertest.cc
@@ -50,7 +50,7 @@
 #include "chrome/browser/profiles/profile_io_data.h"
 #include "chrome/browser/renderer_host/chrome_resource_dispatcher_host_delegate.h"
 #include "chrome/browser/safe_browsing/local_database_manager.h"
-#include "chrome/browser/safe_browsing/safe_browsing_service.h"
+#include "chrome/browser/safe_browsing/test_safe_browsing_service.h"
 #include "chrome/browser/task_management/providers/web_contents/web_contents_tags_manager.h"
 #include "chrome/browser/task_management/task_management_browsertest_util.h"
 #include "chrome/browser/task_manager/task_manager.h"
@@ -741,20 +741,13 @@ class TestPrerenderContentsFactory : public PrerenderContents::Factory {
   std::deque<ExpectedContents> expected_contents_queue_;
 };
 
-// TODO(nparker): Switch this to use TestSafeBrowsingDatabaseManager and run
-// with SAFE_BROWSING_DB_LOCAL || SAFE_BROWSING_DB_REMOTE.
-// Note: don't forget to override GetProtocolManagerDelegate and return NULL,
-// because FakeSafeBrowsingDatabaseManager does not implement
-// LocalSafeBrowsingDatabaseManager.
-#if defined(FULL_SAFE_BROWSING)
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)

[Nathan Parker, 2016/05/07 00:05:32]

Is this file included in Android tests?  If so, then this will now test on
Android when it didn't before, which is good.  And... I think this #if statement
will be true on all platforms (except iOS, which doesn't use /chrome/**), which
could mean you could remove all the ifdefs!

[Jialiu Lin, 2016/05/12 21:53:58]

Good to know. ifdefs removed. Thanks!

 // A SafeBrowsingDatabaseManager implementation that returns a fixed result for
 // a given URL.
 class FakeSafeBrowsingDatabaseManager
-    : public LocalSafeBrowsingDatabaseManager {
+    : public safe_browsing::TestSafeBrowsingDatabaseManager {
  public:
-  explicit FakeSafeBrowsingDatabaseManager(SafeBrowsingService* service)
-      : LocalSafeBrowsingDatabaseManager(service),
-        threat_type_(safe_browsing::SB_THREAT_TYPE_SAFE) {}
+  FakeSafeBrowsingDatabaseManager() {}
 
   // Called on the IO thread to check if the given url is safe or not.  If we
   // can synchronously determine that the url is safe, CheckUrl returns true.
@@ -766,7 +759,8 @@ class FakeSafeBrowsingDatabaseManager
   // client, and false will be returned).
   // Overrides SafeBrowsingDatabaseManager::CheckBrowseUrl.
   bool CheckBrowseUrl(const GURL& gurl, Client* client) override {
-    if (gurl != url_ || threat_type_ == safe_browsing::SB_THREAT_TYPE_SAFE)
+    if (badurls_.find(gurl.spec())==badurls_.end() ||

[Nathan Parker, 2016/05/07 00:05:32]

nit: spaces around ==

[Jialiu Lin, 2016/05/12 21:53:58]

Done.

+        badurls_[gurl.spec()] == safe_browsing::SB_THREAT_TYPE_SAFE)
       return true;
 
     BrowserThread::PostTask(
@@ -777,8 +771,20 @@ class FakeSafeBrowsingDatabaseManager
   }
 
   void SetThreatTypeForUrl(const GURL& url, SBThreatType threat_type) {
-    url_ = url;
-    threat_type_ = threat_type;
+    badurls_[url.spec()] = threat_type;
+  }
+
+  // These are called when checking URLs, so we implement them.
+  bool IsSupported() const override { return true; }
+  bool ChecksAreAlwaysAsync() const override { return false; }
+  bool CanCheckResourceType(
+      content::ResourceType /* resource_type */) const override {
+    return true;
+  }
+
+  bool CheckExtensionIDs(const std::set<std::string>& extension_ids,
+                         Client* client) override {
+    return true;
   }
 
  private:
@@ -796,59 +802,14 @@ class FakeSafeBrowsingDatabaseManager
         client,
         safe_browsing::MALWARE,
         expected_threats);
-    sb_check.url_results[0] = threat_type_;
+    sb_check.url_results[0] = badurls_[gurl.spec()];
     sb_check.OnSafeBrowsingResult();
   }
 
-  GURL url_;
-  SBThreatType threat_type_;
+  base::hash_map<std::string, SBThreatType> badurls_;
   DISALLOW_COPY_AND_ASSIGN(FakeSafeBrowsingDatabaseManager);
 };
 
-class FakeSafeBrowsingService : public SafeBrowsingService {
- public:
-  FakeSafeBrowsingService() { }
-
-  // Returned pointer has the same lifespan as the database_manager_ refcounted
-  // object.
-  FakeSafeBrowsingDatabaseManager* fake_database_manager() {
-    return fake_database_manager_;
-  }
-
- protected:
-  ~FakeSafeBrowsingService() override {}
-
-  safe_browsing::SafeBrowsingDatabaseManager* CreateDatabaseManager() override {
-    fake_database_manager_ = new FakeSafeBrowsingDatabaseManager(this);
-    return fake_database_manager_;
-  }
-
- private:
-  FakeSafeBrowsingDatabaseManager* fake_database_manager_;
-
-  DISALLOW_COPY_AND_ASSIGN(FakeSafeBrowsingService);
-};
-
-// Factory that creates FakeSafeBrowsingService instances.
-class TestSafeBrowsingServiceFactory
-    : public safe_browsing::SafeBrowsingServiceFactory {
- public:
-  TestSafeBrowsingServiceFactory() :
-      most_recent_service_(NULL) { }
-  ~TestSafeBrowsingServiceFactory() override {}
-
-  SafeBrowsingService* CreateSafeBrowsingService() override {
-    most_recent_service_ =  new FakeSafeBrowsingService();
-    return most_recent_service_;
-  }
-
-  FakeSafeBrowsingService* most_recent_service() const {
-    return most_recent_service_;
-  }
-
- private:
-  FakeSafeBrowsingService* most_recent_service_;
-};
 #endif
 
 class FakeDevToolsClient : public content::DevToolsAgentHostClient {
@@ -1130,8 +1091,9 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   PrerenderBrowserTest()
       : autostart_test_server_(true),
         prerender_contents_factory_(NULL),
-#if defined(FULL_SAFE_BROWSING)
-        safe_browsing_factory_(new TestSafeBrowsingServiceFactory()),
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
+        safe_browsing_factory_(
+            new safe_browsing::TestSafeBrowsingServiceFactory()),
 #endif
         call_javascript_(true),
         check_load_events_(true),
@@ -1149,13 +1111,15 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   void SetUpInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
+    safe_browsing_factory_->SetTestDatabaseManager(
+        new FakeSafeBrowsingDatabaseManager());
     SafeBrowsingService::RegisterFactory(safe_browsing_factory_.get());
 #endif
   }
 
   void TearDownInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
     SafeBrowsingService::RegisterFactory(NULL);
 #endif
   }
@@ -1186,6 +1150,12 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
     ASSERT_TRUE(prerender_contents_factory_ == NULL);
     prerender_contents_factory_ = new TestPrerenderContentsFactory;
     prerender_manager->SetPrerenderContentsFactory(prerender_contents_factory_);
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
+    // Note: Need to make GetProtocolManagerDelegate return NULL,
+    // because FakeSafeBrowsingDatabaseManager does not implement
+    // LocalSafeBrowsingDatabaseManager.
+    DCHECK(safe_browsing_factory_->test_safe_browsing_service());

[Nathan Parker, 2016/05/07 00:05:32]

I'm not not sure how the comment applies to this DCHECK.  Is it a todo?

ASSERT_TRUE might be better, so it'll check in non-debug builds as well.

[Jialiu Lin, 2016/05/12 21:53:58]

Oops, sorry it is a todo for myself. Comments removed and changed to Assert_true

+#endif
   }
 
   // Convenience function to get the currently active WebContents in
@@ -1482,10 +1452,11 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
     return static_cast<int>(history_list->GetSize());
   }
 
-#if defined(FULL_SAFE_BROWSING)
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
   FakeSafeBrowsingDatabaseManager* GetFakeSafeBrowsingDatabaseManager() {
-    return safe_browsing_factory_->most_recent_service()->
-        fake_database_manager();
+    return static_cast<FakeSafeBrowsingDatabaseManager*>(
+      safe_browsing_factory_->test_safe_browsing_service()
+      ->database_manager().get());
   }
 #endif
 
@@ -1724,8 +1695,9 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   TestPrerenderContentsFactory* prerender_contents_factory_;
-#if defined(FULL_SAFE_BROWSING)
-  std::unique_ptr<TestSafeBrowsingServiceFactory> safe_browsing_factory_;
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
+  std::unique_ptr<
+      safe_browsing::TestSafeBrowsingServiceFactory> safe_browsing_factory_;
 #endif
   NeverRunsExternalProtocolHandlerDelegate external_protocol_handler_delegate_;
   GURL dest_url_;
@@ -2996,7 +2968,7 @@ IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderSSLClientCertIframe) {
                    0);
 }
 
-#if defined(FULL_SAFE_BROWSING)
+#if defined(SAFE_BROWSING_DB_LOCAL) || defined(SAFE_BROWSING_DB_REMOTE)
 // Ensures that we do not prerender pages with a safe browsing
 // interstitial.
 IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderSafeBrowsingTopLevel) {