diff --git a/llvm/unittests/Analysis/ProfileSummaryInfoTest.cpp b/llvm/unittests/Analysis/ProfileSummaryInfoTest.cpp
--- a/llvm/unittests/Analysis/ProfileSummaryInfoTest.cpp
+++ b/llvm/unittests/Analysis/ProfileSummaryInfoTest.cpp
@@ -240,6 +240,11 @@
   MDBuilder MDB(M->getContext());
   CI2->setMetadata(llvm::LLVMContext::MD_prof, MDB.createBranchWeights({400}));
   EXPECT_FALSE(PSI.isHotCallSite(CS2, &BFI));
+
+  EXPECT_TRUE(PSI.isFunctionHotInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionColdInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionHotInCallGraphNthPercentile(10000, F, BFI));
+  EXPECT_TRUE(PSI.isFunctionColdInCallGraphNthPercentile(10000, F, BFI));
 }
 
 TEST_F(ProfileSummaryInfoTest, InstrProfNoFuncEntryCount) {
@@ -265,6 +270,9 @@
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB1, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB2, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB3, &BFI));
+
+  EXPECT_FALSE(PSI.isFunctionHotInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionColdInCallGraphNthPercentile(990000, F, BFI));
 }
 
 TEST_F(ProfileSummaryInfoTest, SampleProf) {
@@ -332,25 +340,10 @@
   CI2->setMetadata(llvm::LLVMContext::MD_prof, MDB.createBranchWeights({400}));
   EXPECT_TRUE(PSI.isHotCallSite(CS2, &BFI));
 
-  {
-    Function *F = M->getFunction("l");
-    BlockFrequencyInfo BFI = buildBFI(*F);
-    BasicBlock &BB0 = F->getEntryBlock();
-    BasicBlock *BB1 = BB0.getTerminator()->getSuccessor(0);
-    BasicBlock *BB2 = BB0.getTerminator()->getSuccessor(1);
-    BasicBlock *BB3 = BB1->getSingleSuccessor();
-
-    // Without the entry count, all should return false.
-    EXPECT_FALSE(PSI.isHotBlockNthPercentile(990000, &BB0, &BFI));
-    EXPECT_FALSE(PSI.isHotBlockNthPercentile(990000, BB1, &BFI));
-    EXPECT_FALSE(PSI.isHotBlockNthPercentile(990000, BB2, &BFI));
-    EXPECT_FALSE(PSI.isHotBlockNthPercentile(990000, BB3, &BFI));
-
-    EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, &BB0, &BFI));
-    EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB1, &BFI));
-    EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB2, &BFI));
-    EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB3, &BFI));
-  }
+  EXPECT_TRUE(PSI.isFunctionHotInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionColdInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionHotInCallGraphNthPercentile(10000, F, BFI));
+  EXPECT_TRUE(PSI.isFunctionColdInCallGraphNthPercentile(10000, F, BFI));
 }
 
 TEST_F(ProfileSummaryInfoTest, SampleProfNoFuncEntryCount) {
@@ -376,6 +369,9 @@
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB1, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB2, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB3, &BFI));
+
+  EXPECT_FALSE(PSI.isFunctionHotInCallGraphNthPercentile(990000, F, BFI));
+  EXPECT_FALSE(PSI.isFunctionColdInCallGraphNthPercentile(990000, F, BFI));
 }
 
 } // end anonymous namespace