Expose malware/phishing/etc. distinction from SafeBrowsingUIManager

chrome/browser/safe_browsing/ui_manager_unittest.cc

Index: chrome/browser/safe_browsing/ui_manager_unittest.cc
diff --git a/chrome/browser/safe_browsing/ui_manager_unittest.cc b/chrome/browser/safe_browsing/ui_manager_unittest.cc
index 9d08f6bfddc4d5a9db0b3bd07d6d822358d8e550..d0dfd2477deade33d410d378a47f41485b402616 100644
--- a/chrome/browser/safe_browsing/ui_manager_unittest.cc
+++ b/chrome/browser/safe_browsing/ui_manager_unittest.cc
@@ -87,7 +87,7 @@ class SafeBrowsingUIManagerTest : public ChromeRenderViewHostTestHarness {
     ui_manager_->AddToWhitelistUrlSet(
         SafeBrowsingUIManager::GetMainFrameWhitelistUrlForResourceForTesting(
             resource),
-        web_contents(), false);
+        web_contents(), false, resource.threat_type);
   }
 
   SafeBrowsingUIManager::UnsafeResource MakeUnsafeResource(
@@ -148,6 +148,21 @@ TEST_F(SafeBrowsingUIManagerTest, WhitelistIgnoresSitesNotAdded) {
   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);