Expose malware/phishing/etc. distinction from SafeBrowsingUIManager

chrome/browser/safe_browsing/ui_manager_unittest.cc

 // Copyright 2015 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.
 
 #include "chrome/browser/safe_browsing/ui_manager.h"
 
 #include "base/run_loop.h"
 #include "chrome/browser/safe_browsing/safe_browsing_blocking_page.h"
 #include "chrome/browser/safe_browsing/safe_browsing_service.h"
 #include "chrome/browser/safe_browsing/ui_manager.h"
@@ skipping 69 matching lines @@
   void TearDown() override { ChromeRenderViewHostTestHarness::TearDown(); }
 
   bool IsWhitelisted(SafeBrowsingUIManager::UnsafeResource resource) {
     return ui_manager_->IsWhitelisted(resource);
   }
 
   void AddToWhitelist(SafeBrowsingUIManager::UnsafeResource resource) {
     ui_manager_->AddToWhitelistUrlSet(
         SafeBrowsingUIManager::GetMainFrameWhitelistUrlForResourceForTesting(
             resource),
-        web_contents(), false);
+        web_contents(), false, resource.threat_type);
   }
 
   SafeBrowsingUIManager::UnsafeResource MakeUnsafeResource(
       const char* url,
       bool is_subresource) {
     SafeBrowsingUIManager::UnsafeResource resource;
     resource.url = GURL(url);
     resource.is_subresource = is_subresource;
     resource.web_contents_getter =
         SafeBrowsingUIManager::UnsafeResource::GetWebContentsGetter(
@@ skipping 40 matching lines @@
   AddToWhitelist(resource);
   EXPECT_TRUE(IsWhitelisted(resource));
 }
 
 TEST_F(SafeBrowsingUIManagerTest, WhitelistIgnoresSitesNotAdded) {
   SafeBrowsingUIManager::UnsafeResource resource =
       MakeUnsafeResourceAndStartNavigation(kGoodURL);
   EXPECT_FALSE(IsWhitelisted(resource));
 }
 
+TEST_F(SafeBrowsingUIManagerTest, WhitelistRemembersThreatType) {
+  SafeBrowsingUIManager::UnsafeResource resource =
+      MakeUnsafeResourceAndStartNavigation(kBadURL);
+  AddToWhitelist(resource);
+  EXPECT_TRUE(IsWhitelisted(resource));
+  SBThreatType threat_type;
+  content::NavigationEntry* entry =
+      web_contents()->GetController().GetVisibleEntry();
+  ASSERT_TRUE(entry);
+  EXPECT_TRUE(ui_manager()->IsUrlWhitelistedOrPendingForWebContents(
+      resource.url, resource.is_subresource, entry,
+      resource.web_contents_getter.Run(), true, &threat_type));
+  EXPECT_EQ(resource.threat_type, threat_type);
+}
+
 TEST_F(SafeBrowsingUIManagerTest, WhitelistIgnoresPath) {
   SafeBrowsingUIManager::UnsafeResource resource =
       MakeUnsafeResourceAndStartNavigation(kBadURL);
   AddToWhitelist(resource);
   EXPECT_TRUE(IsWhitelisted(resource));
 
   content::WebContentsTester::For(web_contents())->CommitPendingNavigation();
 
   SafeBrowsingUIManager::UnsafeResource resource_path =
       MakeUnsafeResourceAndStartNavigation(kBadURLWithPath);
@@ skipping 228 matching lines @@
   SimulateBlockingPageDone(resources, true);
   EXPECT_TRUE(delegate.visible_security_state_changed());
 
   waiter.WaitForCallback();
   EXPECT_TRUE(waiter.callback_called());
   EXPECT_TRUE(waiter.proceed());
   EXPECT_TRUE(IsWhitelisted(resource));
 }
 
 }  // namespace safe_browsing