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 87d502e8dbdb5ae62964c0a20759b05fd0806240..a1b672c3a2471bcf93f9fbcd8be730c04cbcf7ce 100644
--- a/chrome/browser/prerender/prerender_browsertest.cc
+++ b/chrome/browser/prerender/prerender_browsertest.cc
@@ -4,6 +4,9 @@
 
 #include <stddef.h>
 #include <deque>
+#include <set>
+#include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -50,7 +53,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_management/task_manager_browsertest_util.h"
@@ -732,20 +735,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.
@@ -757,8 +752,10 @@ 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 (bad_urls_.find(gurl.spec()) == bad_urls_.end() ||
+        bad_urls_[gurl.spec()] == safe_browsing::SB_THREAT_TYPE_SAFE) {
       return true;
+    }
 
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
@@ -768,8 +765,20 @@ class FakeSafeBrowsingDatabaseManager
   }
 
   void SetThreatTypeForUrl(const GURL& url, SBThreatType threat_type) {
-    url_ = url;
-    threat_type_ = threat_type;
+    bad_urls_[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:
@@ -787,61 +796,14 @@ class FakeSafeBrowsingDatabaseManager
         client,
         safe_browsing::MALWARE,
         expected_threats);
-    sb_check.url_results[0] = threat_type_;
+    sb_check.url_results[0] = bad_urls_[gurl.spec()];
     sb_check.OnSafeBrowsingResult();
   }
 
-  GURL url_;
-  SBThreatType threat_type_;
+  std::unordered_map<std::string, SBThreatType> bad_urls_;
   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() {}
@@ -1121,14 +1083,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 {}
 
@@ -1140,15 +1100,13 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   void SetUpInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
-    SafeBrowsingService::RegisterFactory(safe_browsing_factory_.get());
-#endif
+    safe_browsing_factory_->SetTestDatabaseManager(
+        new FakeSafeBrowsingDatabaseManager());
+    SafeBrowsingService::RegisterFactory(safe_browsing_factory_);
   }
 
   void TearDownInProcessBrowserTestFixture() override {
-#if defined(FULL_SAFE_BROWSING)
     SafeBrowsingService::RegisterFactory(NULL);
-#endif
   }
 
   void SetUpCommandLine(base::CommandLine* command_line) override {
@@ -1177,6 +1135,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
@@ -1469,12 +1428,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 =
@@ -1711,9 +1670,7 @@ class PrerenderBrowserTest : virtual public InProcessBrowserTest {
   }
 
   TestPrerenderContentsFactory* prerender_contents_factory_;
-#if defined(FULL_SAFE_BROWSING)
-  std::unique_ptr<TestSafeBrowsingServiceFactory> safe_browsing_factory_;
-#endif
+  safe_browsing::TestSafeBrowsingServiceFactory* safe_browsing_factory_;
   NeverRunsExternalProtocolHandlerDelegate external_protocol_handler_delegate_;
   GURL dest_url_;
   std::unique_ptr<net::EmbeddedTestServer> https_src_server_;
@@ -2975,7 +2932,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) {
@@ -3041,8 +2997,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",