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 5efb134167aa3b57e4a0937e8862aaf9769b0bbf..638452830c3cf7b2c64c00e8900c06901772ea58 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/mock_web_contents_task_manager.h"
 #include "chrome/browser/task_management/providers/web_contents/web_contents_tags_manager.h"
 #include "chrome/browser/task_manager/task_manager.h"
@@ -740,20 +740,12 @@ 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)
 // 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.
@@ -765,7 +757,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() ||
+        badurls_[gurl.spec()] == safe_browsing::SB_THREAT_TYPE_SAFE)
       return true;
 
     BrowserThread::PostTask(
@@ -776,8 +769,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:
@@ -795,61 +800,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 {
  public:
   FakeDevToolsClient() {}
@@ -1129,14 +1087,12 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   PrerenderBrowserTest()
       : autostart_test_server_(true),
         prerender_contents_factory_(NULL),
-#if defined(FULL_SAFE_BROWSING)
-        safe_browsing_factory_(new TestSafeBrowsingServiceFactory()),
-#endif
+        safe_browsing_factory_(
+            new safe_browsing::TestSafeBrowsingServiceFactory()),
         call_javascript_(true),
         check_load_events_(true),
         loader_path_("/prerender/prerender_loader.html"),
-        explicitly_set_browser_(NULL) {
-  }
+        explicitly_set_browser_(NULL) {}
 
   ~PrerenderBrowserTest() override {}
 
@@ -1148,15 +1104,13 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   void SetUpInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
+    safe_browsing_factory_->SetTestDatabaseManager(
+        new FakeSafeBrowsingDatabaseManager());
     SafeBrowsingService::RegisterFactory(safe_browsing_factory_.get());
-#endif
   }
 
   void TearDownInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
     SafeBrowsingService::RegisterFactory(NULL);
-#endif
   }
 
   void SetUpCommandLine(base::CommandLine* command_line) override {
@@ -1185,6 +1139,7 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
     ASSERT_TRUE(prerender_contents_factory_ == NULL);
     prerender_contents_factory_ = new TestPrerenderContentsFactory;
     prerender_manager->SetPrerenderContentsFactory(prerender_contents_factory_);
+    ASSERT_TRUE(safe_browsing_factory_->test_safe_browsing_service());
   }
 
   // Convenience function to get the currently active WebContents in
@@ -1477,12 +1432,12 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
     return static_cast<int>(history_list->GetSize());
   }
 
-#if defined(FULL_SAFE_BROWSING)
   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
 
   TestPrerenderContents* GetPrerenderContentsFor(const GURL& url) const {
     PrerenderManager::PrerenderData* prerender_data =
@@ -1719,9 +1674,8 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   TestPrerenderContentsFactory* prerender_contents_factory_;
-#if defined(FULL_SAFE_BROWSING)
-  std::unique_ptr<TestSafeBrowsingServiceFactory> safe_browsing_factory_;
-#endif
+  std::unique_ptr<safe_browsing::TestSafeBrowsingServiceFactory>
+      safe_browsing_factory_;
   NeverRunsExternalProtocolHandlerDelegate external_protocol_handler_delegate_;
   GURL dest_url_;
   std::unique_ptr<net::EmbeddedTestServer> https_src_server_;
@@ -2983,7 +2937,6 @@ IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderSSLClientCertIframe) {
                    0);
 }
 
-#if defined(FULL_SAFE_BROWSING)
 // Ensures that we do not prerender pages with a safe browsing
 // interstitial.
 IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderSafeBrowsingTopLevel) {
@@ -3049,8 +3002,6 @@ IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderSafeBrowsingIframe) {
                    0);
 }
 
-#endif
-
 // Checks that a local storage read will not cause prerender to fail.
 IN_PROC_BROWSER_TEST_F(PrerenderBrowserTest, PrerenderLocalStorageRead) {
   PrerenderTestURL("/prerender/prerender_localstorage_read.html",